diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 6fc0a97838..2aacab80f5 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -566,7 +566,7 @@ class MCPRequestHandler: ) ) - key_access_group_extras = ( + key_access_group_grants = ( await MCPRequestHandler._get_key_access_group_mcp_server_extras( user_api_key_auth ) @@ -577,11 +577,11 @@ class MCPRequestHandler: ######################################################### key_set = set(allowed_mcp_servers_for_key) team_set = set(allowed_mcp_servers_for_team) - extras_set = set(key_access_group_extras) + grants_set = set(key_access_group_grants) - has_lower_level_mcp_restrictions = bool(key_set or team_set or extras_set) + has_lower_level_mcp_restrictions = bool(key_set or team_set or grants_set) - # 1. Team-gated base scope. + # 1. Key/team ceiling. An empty set means "this level does not restrict". if not team_set: base = key_set # no team restriction elif not key_set: @@ -589,9 +589,10 @@ class MCPRequestHandler: else: base = key_set & team_set # both restrict → intersect - # 2. Extend with access-group extras (LIT-3189 — bypasses team - # ceiling, gated by group's assigned_team_ids / assigned_key_ids). - allowed_mcp_servers: List[str] = list(base | extras_set) + # 2. Add the key's access-group grants on top. These are additive: + # attaching a group to the key grants its servers regardless of the + # team ceiling. + allowed_mcp_servers: List[str] = list(base | grants_set) ######################################################### # Check end_user permissions if end_user_id is set @@ -890,52 +891,12 @@ class MCPRequestHandler: ) -> List[str]: """ Resolve the key's unified `access_group_ids` (LiteLLM_AccessGroupTable) to - MCP server IDs, gated by the access group's `assigned_team_ids` / - `assigned_key_ids`. These servers extend the team's MCP scope rather - than being capped by it. Tag-style `mcp_access_groups` (per-server tags) - are intentionally not handled here — they have no assignment fields and - remain subject to the team ceiling. - """ - if user_api_key_auth is None: - return [] - try: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy.auth.auth_checks import ( - get_authorized_resources_from_key_access_groups, - ) - - raw_server_ids = await get_authorized_resources_from_key_access_groups( - valid_token=user_api_key_auth, - team_object=None, - resource_field="access_mcp_server_ids", - ) - if not raw_server_ids: - return [] - # Permission entries may be server_ids OR names/aliases — expand to ids. - return global_mcp_server_manager.expand_permission_list(raw_server_ids) - except Exception as e: - verbose_logger.warning( - f"Failed to get key access group MCP server extras: {str(e)}" - ) - return [] - - @staticmethod - async def _get_allowed_mcp_servers_for_key( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: - """ - Get allowed MCP servers for a key (the key's own scope). - - Unions two sources: - - Legacy key.object_permission (mcp_servers, mcp_access_groups, - mcp_tool_permissions). - - Unified key.access_group_ids → access_group.access_mcp_server_ids. - Mirrors the ungated fallback in can_key_call_model — the group is - attached to the key itself, so it grants the key's own scope (no - assigned_key_ids re-check). The gated, team-ceiling-busting override - lives in _get_key_access_group_mcp_server_extras. + MCP server IDs as additive grants: a group attached to the key extends the + key's allowed servers on top of the key/team ceiling rather than being + capped by the team. Attaching the group to the key is itself the grant — + no `assigned_key_ids` / `assigned_team_ids` re-check. Tag-style + `mcp_access_groups` (per-server tags) live in the key's object_permission + scope, not here. """ if user_api_key_auth is None: return [] @@ -945,7 +906,6 @@ class MCPRequestHandler: ) from litellm.proxy.auth.auth_checks import ( _get_mcp_server_ids_from_access_groups, - get_object_permission, ) from litellm.proxy.proxy_server import ( prisma_client, @@ -953,17 +913,48 @@ class MCPRequestHandler: user_api_key_cache, ) - # Unified key.access_group_ids → MCP servers (ungated: the group is - # attached to the key, so it grants the key's own scope). Entries in - # access_mcp_server_ids may be server_ids OR names/aliases, so expand - # to ids here — matching the legacy object_permission path below. - key_access_group_servers = global_mcp_server_manager.expand_permission_list( - await _get_mcp_server_ids_from_access_groups( - access_group_ids=user_api_key_auth.access_group_ids or [], - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + raw_server_ids = await _get_mcp_server_ids_from_access_groups( + access_group_ids=user_api_key_auth.access_group_ids or [], + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if not raw_server_ids: + return [] + # Permission entries may be server_ids OR names/aliases — expand to ids. + return global_mcp_server_manager.expand_permission_list(raw_server_ids) + except Exception as e: + verbose_logger.warning( + f"Failed to get key access group MCP server grants: {str(e)}" + ) + return [] + + @staticmethod + async def _get_allowed_mcp_servers_for_key( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[str]: + """ + Get the key's own MCP ceiling from its object_permission + (mcp_servers, tag-style mcp_access_groups, mcp_tool_permissions). + + Unified key.access_group_ids are NOT resolved here — they are additive + grants handled by _get_key_access_group_mcp_server_extras and unioned on + top of the key/team ceiling, so they must not enter this scope (which is + intersected against the team). + """ + if user_api_key_auth is None: + return [] + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.auth.auth_checks import ( + get_object_permission, + ) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, ) # Get key object permission (already loaded in main auth flow, or fetch from DB) @@ -983,7 +974,7 @@ class MCPRequestHandler: proxy_logging_obj=proxy_logging_obj, ) if key_object_permission is None: - return list(set(key_access_group_servers)) + return [] # Permission entries may be server_ids OR names/aliases — expand to ids. direct_mcp_servers = global_mcp_server_manager.expand_permission_list( @@ -1005,12 +996,7 @@ class MCPRequestHandler: ) # Combine all lists - all_servers = ( - direct_mcp_servers - + access_group_servers - + tool_perm_servers - + key_access_group_servers - ) + all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers return list(set(all_servers)) except Exception as e: verbose_logger.warning( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 95294f4838..751f855ea3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -242,6 +242,11 @@ class KeyManagementRoutes(str, enum.Enum): TEAM_KEY_BULK_UPDATE = "/team/key/bulk_update" KEY_RESET_SPEND = "/key/{key_id}/reset_spend" + # Field-level opt-in permission (not a real HTTP route). When present in a + # team's `team_member_permissions`, non-admin members of that team may set + # `access_group_ids` on keys they create/update. Default-deny. + KEY_ACCESS_GROUP_ASSIGNMENT = "/key/access_group_assignment" + # info and health routes KEY_INFO = "/key/info" KEY_HEALTH = "/key/health" @@ -552,6 +557,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.SPEND_LOGS_V2.value, KeyManagementRoutes.KEY_RESET_SPEND.value, KeyManagementRoutes.KEY_ALIASES.value, + KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value, ] management_routes = ( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f49fb1f37b..0e645013b9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -325,6 +325,14 @@ def _team_key_generation_check( _team_key_generation.get("required_params"), ) + # Field-level opt-in: non-admin members may only assign access groups when + # the team has enabled KEY_ACCESS_GROUP_ASSIGNMENT. + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=user_api_key_dict, + team_table=team_table, + access_group_ids=data.access_group_ids, + ) + return True @@ -2267,6 +2275,14 @@ async def _validate_update_key_data( detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.", ) + # Field-level opt-in: non-admin members may only assign access groups when + # the team has enabled KEY_ACCESS_GROUP_ASSIGNMENT. + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=user_api_key_dict, + team_table=team_obj, + access_group_ids=data.access_group_ids, + ) + if team_obj is not None: await _check_team_key_limits( team_table=team_obj, @@ -4511,6 +4527,23 @@ async def regenerate_key_fn( # noqa: PLR0915 detail={"error": "You are not authorized to regenerate this key"}, ) + # Gate access_group_ids on regenerate, same as /key/generate and + # /key/update. Use the existing key's team since the body may omit it. + if data is not None and data.access_group_ids: + regenerate_team_table: Optional[LiteLLM_TeamTableCachedObj] = None + if _key_in_db.team_id is not None: + regenerate_team_table = await get_team_object( + team_id=_key_in_db.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=user_api_key_dict, + team_table=regenerate_team_table, + access_group_ids=data.access_group_ids, + ) + verbose_proxy_logger.info( "Key regeneration requested: key_alias=%s", getattr(_key_in_db, "key_alias", None), diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 50339210a6..2272a37488 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -154,6 +154,70 @@ class TeamMemberPermissionChecks: return True + @staticmethod + def enforce_member_can_assign_access_groups( + user_api_key_dict: UserAPIKeyAuth, + team_table: Optional[LiteLLM_TeamTableCachedObj], + access_group_ids: Optional[List[str]], + ) -> None: + """ + Field-level opt-in gate: a non-admin team member may only set + `access_group_ids` on a (team) key if their team has opted in by adding + `KEY_ACCESS_GROUP_ASSIGNMENT` to `team_member_permissions`. + + Bypassed for proxy admins, team admins, and personal (non-team) keys. + Default-deny: members cannot self-assign access groups until enabled. + + Raises HTTPException(403) when a gated member attempts the assignment. + """ + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _get_user_in_team, + ) + + # No-op when the request does not assign any access groups. + if not access_group_ids: + return + + # Proxy admins always bypass. + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + + # Personal (non-team) keys are out of scope for team-member gating. + if team_table is None: + return + + team_member_object = _get_user_in_team( + team_table=team_table, user_id=user_api_key_dict.user_id + ) + + # Team admins always bypass (consistent with other member-permission checks). + if team_member_object is not None and team_member_object.role == "admin": + return + + permissions = ( + TeamMemberPermissionChecks._get_list_of_route_enum_as_str( + TeamMemberPermissionChecks.get_permissions_for_team_member( + team_member_object=team_member_object, + team_table=team_table, + ) + ) + if team_member_object is not None + else [] + ) + + if KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value not in permissions: + raise HTTPException( + status_code=403, + detail=( + "Team members cannot assign access groups to keys for team " + f"{team_table.team_id}. Ask a team or proxy admin to enable the " + f"'{KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value}' team " + "member permission to allow this." + ), + ) + @staticmethod async def user_belongs_to_keys_team( user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index b3508e1312..95c826daa8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -200,6 +200,117 @@ class TestMCPRequestHandler: result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) assert result == [] # Should handle exception gracefully + @pytest.mark.parametrize( + "key_servers,team_servers,grant_servers,expected,scenario", + [ + # Key has no own scope, restrictive team ceiling {test}, server + # granted only via key.access_group_ids → caller sees team's server + # AND the grant (grant is added on top of the ceiling). + ( + [], + ["test"], + ["context7"], + ["context7", "test"], + "grant_over_team_ceiling", + ), + # key {a} ∩ team {b} = {} ; the grant still surfaces, proving grants + # are unioned with the ceiling, not intersected against it. + ( + ["a"], + ["b"], + ["context7"], + ["context7"], + "grant_survives_empty_intersection", + ), + # No grant → ceiling behavior is unchanged (no additive leakage). + (["x", "y"], ["x"], [], ["x"], "no_grant_keeps_intersection"), + ], + ) + async def test_access_group_grants_are_additive_over_ceiling( + self, key_servers, team_servers, grant_servers, expected, scenario + ): + """Regression: key.access_group_ids grants are unioned on top of the + key/team MCP ceiling, so a grant reaches the caller even when the team + ceiling does not include it (and even when key ∩ team is empty).""" + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + access_group_ids=["grp-mcp"], + ) + with ( + patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key, + patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_team" + ) as mock_team, + patch.object( + MCPRequestHandler, "_get_key_access_group_mcp_server_extras" + ) as mock_grants, + ): + mock_key.return_value = key_servers + mock_team.return_value = team_servers + mock_grants.return_value = grant_servers + result = await MCPRequestHandler.get_allowed_mcp_servers(mock_user_auth) + assert sorted(result) == sorted(expected) + + async def test_access_group_extras_returns_empty_when_no_auth(self): + """No auth object → no additive grants.""" + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras(None) + assert result == [] + + async def test_access_group_extras_returns_empty_without_access_group_ids(self): + """A key with no resolvable access groups yields no additive grants + (the `if not raw_server_ids: return []` branch).""" + auth = UserAPIKeyAuth(api_key="k", access_group_ids=[]) + with ( + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new=AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + auth + ) + assert result == [] + # expand_permission_list must not be reached when there are no raw ids. + mock_mgr.expand_permission_list.assert_not_called() + + async def test_access_group_extras_expands_resolved_server_ids(self): + """Resolved access-group server ids/names are expanded to server ids.""" + auth = UserAPIKeyAuth(api_key="k", access_group_ids=["grp-mcp"]) + with ( + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new=AsyncMock(return_value=["alias-a", "srv-b"]), + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.expand_permission_list.return_value = ["srv-a", "srv-b"] + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + auth + ) + assert sorted(result) == ["srv-a", "srv-b"] + mock_mgr.expand_permission_list.assert_called_once_with(["alias-a", "srv-b"]) + + async def test_access_group_extras_swallows_errors(self): + """Resolution failures degrade to no grants rather than raising.""" + auth = UserAPIKeyAuth(api_key="k", access_group_ids=["grp-mcp"]) + with patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new=AsyncMock(side_effect=Exception("db down")), + ): + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + auth + ) + assert result == [] + @pytest.mark.parametrize( "headers,expected_api_key,expected_mcp_auth_header,expected_server_auth_headers", [ @@ -3335,12 +3446,12 @@ async def test_mcp_key_access_group_extras_when_group_has_no_servers(): @pytest.mark.asyncio -async def test_mcp_key_access_group_extras_when_group_authorizes_neither(): - """ - Escalation regression: team member attaches a foreign access group to their key. - Group grants servers BUT assigned_team_ids/assigned_key_ids exclude this caller. - No extras contributed. - """ +async def test_mcp_key_access_group_extras_granted_even_when_group_authorizes_neither(): + """Grants are ungated: attaching the group to the key is itself the grant, so its + servers are contributed even when assigned_team_ids/assigned_key_ids exclude this + caller. (A team member self-assigning a foreign group to reach past the team + ceiling is a known, accepted-for-now tradeoff; restricting who may set + key.access_group_ids is a separate concern.)""" valid_token = UserAPIKeyAuth( token="team-a-token", access_group_ids=["team-b-mcp-group"], @@ -3365,7 +3476,7 @@ async def test_mcp_key_access_group_extras_when_group_authorizes_neither(): result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( valid_token ) - assert result == [] + assert result == ["srv-finance-only"] finally: _stop_patches(patches) @@ -3650,11 +3761,12 @@ async def test_get_allowed_mcp_servers_includes_team_access_group_extras_end_to_ @pytest.mark.asyncio -async def test_key_access_group_ids_resolves_mcp_servers_ungated(): - """A teamless key whose unified access_group_ids grant an MCP server sees it - even though the group lists the key in NEITHER assigned_key_ids NOR - assigned_team_ids — the group is attached to the key, so it grants the key's - own scope (ungated, mirroring can_key_call_model's fallback).""" +async def test_allowed_mcp_servers_for_key_excludes_access_group_ids(): + """The key's own ceiling (which is intersected against the team) must NOT resolve + access_group_ids — those are additive grants handled separately, so folding them + in here is exactly the bug this fix removes. A key with only access_group_ids and + no object_permission yields an empty ceiling, and the group resolver is never + called from this path.""" auth = UserAPIKeyAuth( token="test-token-hash", api_key="sk-test", @@ -3671,15 +3783,16 @@ async def test_key_access_group_ids_resolves_mcp_servers_ungated(): ): result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth) - assert result == ["srv-stripe"] - mock_resolver.assert_called_once() - assert mock_resolver.call_args.kwargs["access_group_ids"] == ["mcp-premium"] + assert result == [] + mock_resolver.assert_not_called() @pytest.mark.asyncio -async def test_key_access_group_ids_union_with_object_permission(): - """When both legacy key.object_permission and unified key.access_group_ids - grant MCP servers, the final list is their union.""" +async def test_allowed_mcp_servers_for_key_uses_object_permission_not_access_groups(): + """The key's own ceiling is built from object_permission alone. Even when the key + also carries access_group_ids that would resolve to other servers, those grants do + NOT enter this (intersected) scope — only the object_permission server comes back. + """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -3714,82 +3827,40 @@ async def test_key_access_group_ids_union_with_object_permission(): "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", new_callable=AsyncMock, return_value=["srv-stripe"], - ), + ) as mock_resolver, ): result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth) - assert set(result) == {"srv-direct", "srv-stripe"} + assert set(result) == {"srv-direct"} + mock_resolver.assert_not_called() finally: global_mcp_server_manager.registry.pop("srv-direct", None) @pytest.mark.asyncio -async def test_key_access_group_ids_empty_returns_no_extras(): - """Empty key.access_group_ids and no object_permission → resolver called with - [], short-circuits without DB access, returns [].""" - auth = UserAPIKeyAuth( - token="test-token-hash", - api_key="sk-test", - access_group_ids=[], - ) - - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch( - "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", - new_callable=AsyncMock, - return_value=[], - ) as mock_resolver, - ): - result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth) - - assert result == [] - mock_resolver.assert_called_once() - assert mock_resolver.call_args.kwargs["access_group_ids"] == [] - - -@pytest.mark.asyncio -async def test_get_allowed_mcp_servers_key_access_group_base_end_to_end(): - """End-to-end bug repro: a teamless key has an MCP-granting access group on - its access_group_ids, but the group lists the key in NEITHER assigned_key_ids - NOR assigned_team_ids. The gated extras path returns [] (no override), yet the - ungated base key path grants the server → the key sees it through - get_allowed_mcp_servers.""" +async def test_get_allowed_mcp_servers_surfaces_ungated_key_access_group_grant_end_to_end(): + """End-to-end: a teamless key has an MCP-granting access group on its + access_group_ids. The grant is resolved ungated by the additive extras path and + surfaces through get_allowed_mcp_servers, even though the key's own ceiling + (object_permission) is empty.""" auth = UserAPIKeyAuth( token="test-token", api_key="sk-test", access_group_ids=["mcp-group"], ) - # Group grants the server but admits neither this key nor its (absent) team. - fake_ag = _fake_mcp_access_group( - access_group_id="mcp-group", - access_mcp_server_ids=["srv-deepwiki"], - assigned_team_ids=[], - assigned_key_ids=[], - ) patches = _patch_proxy_server_globals_for_mcp() + [ - # Ungated base resolver used by _get_allowed_mcp_servers_for_key. patch( "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", new_callable=AsyncMock, return_value=["srv-deepwiki"], ), - # Gated path (_get_key_access_group_mcp_server_extras) resolves the group - # via get_access_object; empty assigned_* → it contributes nothing. - patch( - "litellm.proxy.auth.auth_checks.get_access_object", - new_callable=AsyncMock, - return_value=fake_ag, - ), ] _start_patches(patches) try: - # Sanity: the gated extras path alone denies (the old behavior). extras = await MCPRequestHandler._get_key_access_group_mcp_server_extras(auth) - assert extras == [] + assert extras == ["srv-deepwiki"] - # But the key now sees the server via the ungated base path. result = await MCPRequestHandler.get_allowed_mcp_servers(auth) assert result == ["srv-deepwiki"] finally: diff --git a/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py b/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py index 6eb05aaf9d..29aa75a0f0 100644 --- a/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py +++ b/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py @@ -265,3 +265,112 @@ class TestCanTeamMemberExecuteKeyManagementEndpoint: user_api_key_cache=MagicMock(), existing_key_row=existing_key_row, ) + + +class TestEnforceMemberCanAssignAccessGroups: + """Opt-in gate controlling whether a non-admin team member may set + `access_group_ids` on a key (generate/update/regenerate).""" + + AG_PERMISSION = KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value + + def _user(self, role="internal_user", user_id="user-a"): + u = MagicMock() + u.user_role = role + u.user_id = user_id + return u + + def _team(self, team_member_permissions, team_id="team-a"): + team = MagicMock() + team.team_id = team_id + team.team_member_permissions = team_member_permissions + return team + + def test_no_access_group_ids_is_noop(self, monkeypatch): + """When no access groups are requested the gate never raises, even + for a gated member with no opt-in permission.""" + from litellm.proxy.management_endpoints import key_management_endpoints + + monkeypatch.setattr( + key_management_endpoints, + "_get_user_in_team", + lambda **kwargs: Member(role="user", user_id="user-a"), + ) + + # Both None and empty list are no-ops. + for access_group_ids in (None, []): + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=self._user(), + team_table=self._team([]), + access_group_ids=access_group_ids, + ) + + def test_proxy_admin_bypasses(self, monkeypatch): + """Proxy admins may assign access groups regardless of team opt-in.""" + from litellm.proxy._types import LitellmUserRoles + + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=self._user(role=LitellmUserRoles.PROXY_ADMIN.value), + team_table=self._team([]), + access_group_ids=["ag-1"], + ) + + def test_personal_key_out_of_scope(self): + """Personal (non-team) keys are not gated by team-member permissions.""" + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=self._user(), + team_table=None, + access_group_ids=["ag-1"], + ) + + def test_team_admin_bypasses(self, monkeypatch): + """Team admins may assign access groups even without the opt-in perm.""" + from litellm.proxy.management_endpoints import key_management_endpoints + + monkeypatch.setattr( + key_management_endpoints, + "_get_user_in_team", + lambda **kwargs: Member(role="admin", user_id="user-a"), + ) + + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=self._user(), + team_table=self._team([]), + access_group_ids=["ag-1"], + ) + + def test_member_denied_without_opt_in(self, monkeypatch): + """A non-admin member without the opt-in permission gets a 403.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints import key_management_endpoints + + monkeypatch.setattr( + key_management_endpoints, + "_get_user_in_team", + lambda **kwargs: Member(role="user", user_id="user-a"), + ) + + with pytest.raises(HTTPException) as exc: + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=self._user(), + team_table=self._team(["/key/generate", "/key/update"]), + access_group_ids=["ag-1"], + ) + assert exc.value.status_code == 403 + assert self.AG_PERMISSION in str(exc.value.detail) + + def test_member_allowed_with_opt_in(self, monkeypatch): + """A non-admin member is allowed once the team opts in via the perm.""" + from litellm.proxy.management_endpoints import key_management_endpoints + + monkeypatch.setattr( + key_management_endpoints, + "_get_user_in_team", + lambda **kwargs: Member(role="user", user_id="user-a"), + ) + + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( + user_api_key_dict=self._user(), + team_table=self._team(["/key/generate", self.AG_PERMISSION]), + access_group_ids=["ag-1"], + ) diff --git a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx index 4a48baec32..bcd910e377 100644 --- a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx +++ b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx @@ -20,6 +20,8 @@ export const PERMISSION_DESCRIPTIONS: Record = { "/key/list": "Member can list virtual keys belonging to this team", "/key/block": "Member can block a virtual key belonging to this team", "/key/unblock": "Member can unblock a virtual key belonging to this team", + "/key/access_group_assignment": + "Member can assign access groups to virtual keys for this team", "/team/daily/activity": "Member can view all team usage data (not just their own)", "/spend/logs":