Merge pull request #40196 from BerriAI/litellm_team_admin_all_proxy_models

fix(ui): let team admins grant a team all proxy models
This commit is contained in:
ryan-crabbe-berri 2026-09-14 15:05:01 -07:00 committed by GitHub
commit 7b6d223105
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 342 additions and 69 deletions

View file

@ -4418,6 +4418,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
access_group_mcp_server_ids: list[str] | None = None
access_group_agent_ids: list[str] | None = None
access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
# Parent org's model ceiling, reported only to callers who can manage the team.
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
organization_models: list[str] | None = None
class TeamInfoResponseObject(TypedDict):

View file

@ -431,27 +431,26 @@ async def _refresh_cached_team(
)
async def _can_manage_team(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return True
return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
async def _verify_team_access(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Verify the caller is authorized to manage the given team.
Access is granted if:
- Caller is a proxy admin, OR
- Caller is an org admin for the team's organization, OR
- Caller is a team admin of this team
Raises HTTPException(403) otherwise.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
"""Raise HTTPException(403) unless the caller can manage the given team."""
if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
return
raise HTTPException(
@ -4370,6 +4369,20 @@ async def _hydrate_member_user_details(
return tuple(hydrate(m) for m in members)
class _OrganizationModelsRow(BaseModel):
models: list[str] = [] # mutable-ok: pydantic field default
class _TeamRowWithOrganization(BaseModel):
litellm_organization_table: _OrganizationModelsRow | None = None
def _parent_organization_models(team_row: BaseModel) -> list[str] | None:
"""Return the parent org's model allow-list, or None when the team has no org."""
organization: Final = _TeamRowWithOrganization.model_validate(team_row.model_dump()).litellm_organization_table
return organization.models if organization is not None else None
async def _resolve_team_access_group_resources(
_team_info: TeamInfoResponseObjectTeamTable,
) -> TeamInfoResponseObjectTeamTable:
@ -4441,7 +4454,11 @@ async def team_info(
try:
team_info: BaseModel | None = await _team_db(prisma_client).find_unique(
where={"team_id": team_id},
include={"litellm_model_table": True, "object_permission": True},
include={
"litellm_model_table": True,
"object_permission": True,
"litellm_organization_table": True,
},
)
if team_info is None:
raise Exception
@ -4450,9 +4467,12 @@ async def team_info(
status_code=status.HTTP_404_NOT_FOUND,
detail={"message": f"Team not found, passed team id: {team_id}."},
)
await validate_membership(
user_api_key_dict=user_api_key_dict,
team_table=LiteLLM_TeamTable.model_validate(team_info.model_dump()),
team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
organization_models: Final[list[str] | None] = (
_parent_organization_models(team_info)
if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
else None
)
## GET ALL KEYS ##
@ -4512,7 +4532,10 @@ async def team_info(
members=resolved_team_info.members_with_roles,
)
hydrated_team_info: Final = resolved_team_info.model_copy(
update={"members_with_roles": hydrated_members} # mutable-ok: pydantic update payload
update={ # mutable-ok: pydantic update payload
"members_with_roles": hydrated_members,
"organization_models": organization_models,
}
)
response_object: Final = TeamInfoResponseObject(

View file

@ -14366,3 +14366,123 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids
assert exc_info.value.status_code == 400
assert expected_error in str(exc_info.value.detail)
mock_db_client.db.query_raw.assert_not_called()
class _TeamRowWithOrganization(LiteLLM_TeamTable):
litellm_organization_table: LiteLLM_OrganizationTable | None = None
@pytest.mark.parametrize(
"organization, expected_models",
[
(
LiteLLM_OrganizationTable(
organization_id="org-1",
budget_id="budget-1",
models=["all-proxy-models"],
created_by="admin",
updated_by="admin",
),
["all-proxy-models"],
),
(
LiteLLM_OrganizationTable(
organization_id="org-1",
budget_id="budget-1",
models=["gpt-4o"],
created_by="admin",
updated_by="admin",
),
["gpt-4o"],
),
(None, None),
],
)
@pytest.mark.asyncio
async def test_team_info_returns_parent_organization_models(organization, expected_models):
"""/team/info must report the parent org's model ceiling.
A team admin who is not an org admin gets a 403 from /organization/info, so this
is the only read that can tell the Admin UI whether the org allows all proxy
models. Without it the team edit form hides the "All Proxy Models" option and a
team admin cannot grant their team everything on the proxy.
"""
from fastapi import Request
from litellm.proxy.management_endpoints import team_endpoints
team_row = _TeamRowWithOrganization(
team_id="team-1",
organization_id="org-1" if organization is not None else None,
litellm_organization_table=organization,
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
mock_prisma.get_data = AsyncMock(return_value=[])
memberships = AsyncMock(return_value=[])
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: no seam on team_info
patch.object(team_endpoints, "get_all_team_memberships", memberships), # test-quality-ok: no seam on team_info
):
response = await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="team-1",
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert response["team_info"].organization_models == expected_models
@pytest.mark.parametrize(
"caller, expected_models",
[
(UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.INTERNAL_USER), ["gpt-4o"]),
(UserAPIKeyAuth(user_id="member-1", user_role=LitellmUserRoles.INTERNAL_USER), None),
(UserAPIKeyAuth(team_id="team-1"), None),
],
)
@pytest.mark.asyncio
async def test_team_info_reports_parent_organization_models_only_to_team_managers(caller, expected_models):
"""Plain members and team keys can read their team, but not the org's wider allow-list."""
from fastapi import Request
from litellm.proxy.management_endpoints import team_endpoints
team_row = _TeamRowWithOrganization(
team_id="team-1",
organization_id="org-1",
members_with_roles=[
Member(user_id="admin-1", role="admin"),
Member(user_id="member-1", role="user"),
],
litellm_organization_table=LiteLLM_OrganizationTable(
organization_id="org-1",
budget_id="budget-1",
models=["gpt-4o"],
created_by="admin",
updated_by="admin",
),
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma.get_data = AsyncMock(return_value=[])
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: no seam on team_info
patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), # test-quality-ok: no seam on team_info
patch.object( # test-quality-ok: no seam on team_info
team_endpoints, "_is_user_org_admin_for_team", AsyncMock(return_value=False)
),
):
response = await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="team-1",
user_api_key_dict=caller,
)
assert response["team_info"].organization_models == expected_models

View file

@ -328,7 +328,8 @@ describe("useTeam", () => {
});
it("should return team data when query is successful", async () => {
(teamInfoCall as any).mockResolvedValue(mockTeams[0]);
// /team/info answers with an envelope; the hook is typed as the team itself.
(teamInfoCall as any).mockResolvedValue({ team_id: "team-1", team_info: mockTeams[0], keys: [] });
const { result } = renderHook(() => useTeam("team-1"), { wrapper });

View file

@ -163,7 +163,8 @@ export const useTeam = (teamId?: string) => {
throw new Error("Missing auth or teamId");
}
return teamInfoCall(accessToken, teamId);
const { team_info } = (await teamInfoCall(accessToken, teamId)) as { team_info: Team };
return team_info;
},
initialData: () => {

View file

@ -242,13 +242,11 @@ describe("ProjectDetail", () => {
it("should show team information when team data is available", () => {
mockUseTeam.mockReturnValue({
data: {
team_info: {
team_id: "team-1",
team_alias: "Engineering",
models: ["gpt-4"],
spend: 50,
members_with_roles: [],
},
team_id: "team-1",
team_alias: "Engineering",
models: ["gpt-4"],
spend: 50,
members_with_roles: [],
},
isLoading: false,
});

View file

@ -14,16 +14,6 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { EditProjectModal } from "./ProjectModals/EditProjectModal";
import { ProjectKeysSection } from "./ProjectKeysSection";
interface TeamInfoShape {
team_id: string;
team_alias?: string;
models?: string[];
max_budget?: number | null;
budget_duration?: string | null;
spend?: number;
members_with_roles?: { user_id: string; role: string }[];
}
interface ProjectDetailProps {
projectId: string;
onBack: () => void;
@ -33,10 +23,7 @@ const utilisationTone = (percent: number) => (percent >= 90 ? "over" : percent >
export function ProjectDetail({ projectId, onBack }: ProjectDetailProps) {
const { data: project, isLoading } = useProjectDetails(projectId);
const { data: teamData } = useTeam(project?.team_id ?? undefined);
// teamInfoCall returns { team_id, team_info: {...}, keys, team_memberships }
const teamInfo: TeamInfoShape | undefined = ((teamData as unknown as { team_info?: TeamInfoShape })?.team_info ??
teamData) as TeamInfoShape | undefined;
const { data: teamInfo } = useTeam(project?.team_id ?? undefined);
const [isEditModalVisible, setIsEditModalVisible] = useState(false);
const spend = project?.spend ?? 0;

View file

@ -415,6 +415,139 @@ describe("ModelSelect", () => {
}
});
it("should take the org model ceiling from the team when /organization/info is not readable", async () => {
const testCases = [
{
name: "org allows all proxy models",
organizationModels: ["all-proxy-models"],
shouldShowSentinel: true,
offered: ["gpt-4", "claude-3"],
notOffered: [] as string[],
},
{
name: "org places no ceiling at all",
organizationModels: [],
shouldShowSentinel: true,
offered: ["gpt-4", "claude-3"],
notOffered: [] as string[],
},
{
name: "org restricts the team to one model",
organizationModels: ["gpt-4"],
shouldShowSentinel: false,
offered: ["gpt-4"],
notOffered: ["claude-3"],
},
];
for (const testCase of testCases) {
const user = userEvent.setup();
// A team admin gets a 403 from /organization/info, so the org query never resolves.
mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
mockUseTeam.mockReturnValue({
data: { team_id: "team-1", organization_models: testCase.organizationModels },
isLoading: false,
} as any);
const { unmount } = renderWithProviders(
<ModelSelect
onChange={mockOnChange}
context="team"
teamID="team-1"
organizationID="org-1"
options={{ includeSpecialOptions: true }}
/>,
);
await openModelList(user);
if (testCase.shouldShowSentinel) {
expectOffered("All Proxy Models");
} else {
expectNotOffered("All Proxy Models");
}
expectOffered("No Default Models");
testCase.offered.forEach(expectOffered);
testCase.notOffered.forEach(expectNotOffered);
unmount();
}
});
it("should stay in the loading state while a list-seeded team is still fetching its org ceiling", () => {
mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
mockUseTeam.mockReturnValue({
data: { team_id: "team-1", models: [] },
isLoading: false,
isFetching: true,
} as any);
renderWithProviders(
<ModelSelect
onChange={mockOnChange}
context="team"
teamID="team-1"
organizationID="org-1"
options={{ includeSpecialOptions: true }}
/>,
);
expect(screen.queryAllByRole("combobox")).toHaveLength(0);
});
it("should not hold the loading state on a background refetch once the org ceiling is known", async () => {
const user = userEvent.setup();
mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
mockUseTeam.mockReturnValue({
data: { team_id: "team-1", organization_models: ["all-proxy-models"] },
isLoading: false,
isFetching: true,
} as any);
renderWithProviders(
<ModelSelect
onChange={mockOnChange}
context="team"
teamID="team-1"
organizationID="org-1"
options={{ includeSpecialOptions: true }}
/>,
);
await openModelList(user);
expectOffered("All Proxy Models");
});
it("should offer no models for an org team when neither the team nor the org reports a ceiling", async () => {
const testCases = [
{ name: "/team/info withheld the ceiling", team: { team_id: "team-1", organization_models: null } },
{ name: "/team/info failed after the list seeded the team", team: { team_id: "team-1", models: [] } },
];
for (const testCase of testCases) {
const user = userEvent.setup();
mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any);
mockUseTeam.mockReturnValue({ data: testCase.team, isLoading: false, isFetching: false } as any);
const { unmount } = renderWithProviders(
<ModelSelect
onChange={mockOnChange}
context="team"
teamID="team-1"
organizationID="org-1"
options={{ includeSpecialOptions: true }}
/>,
);
await openModelList(user);
expectNotOffered("All Proxy Models");
expectOffered("No Default Models");
expectNotOffered("gpt-4");
expectNotOffered("claude-3");
unmount();
}
});
it("should use custom dataTestId when provided", async () => {
renderWithProviders(
<ModelSelect

View file

@ -19,7 +19,7 @@ import {
} from "@/components/ui/combobox";
import { Skeleton } from "@/components/ui/skeleton";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { Organization, Team } from "../networking";
import type { Team } from "@/components/key_team_helpers/key_list";
import { splitWildcardModels } from "./modelUtils";
const MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE = {
@ -69,12 +69,19 @@ type ModelOptionGroup = {
type FilterContextArgs = {
allProxyModels: string[];
selectedTeam?: Team;
selectedOrganization?: Organization;
organizationID?: string;
organizationModels?: string[];
userModels?: string[];
options?: ModelSelectProps["options"];
};
const isUncappedModelCeiling = (organizationModels: string[]) =>
organizationModels.length === 0 || organizationModels.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value);
// useTeam seeds from the team list, which omits organization_models; /team/info is the only source of the org ceiling.
const isAwaitingOrganizationModels = (team: Team | undefined, isFetchingTeam: boolean) =>
isFetchingTeam && team !== undefined && team.organization_models === undefined;
const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextArgs) => string[]> = {
user: ({ allProxyModels, userModels, options }) => {
if (!userModels) return [];
@ -82,18 +89,10 @@ const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextAr
return [];
},
team: ({ allProxyModels, selectedOrganization, userModels }) => {
if (selectedOrganization) {
if (
selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) ||
selectedOrganization.models.length === 0
) {
return allProxyModels;
}
return allProxyModels.filter((model) => selectedOrganization.models.includes(model));
}
return allProxyModels ?? [];
team: ({ allProxyModels, organizationID, organizationModels }) => {
if (organizationModels === undefined) return organizationID ? [] : allProxyModels;
if (isUncappedModelCeiling(organizationModels)) return allProxyModels;
return allProxyModels.filter((model) => organizationModels.includes(model));
},
organization: ({ allProxyModels }) => {
@ -108,7 +107,7 @@ const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextAr
const filterModels = (
allProxyModels: ProxyModel[],
ctx: ModelSelectProps,
extra: { selectedTeam?: Team; selectedOrganization?: Organization; userModels?: string[] },
extra: { organizationModels?: string[]; userModels?: string[] },
): string[] => {
const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((m) => [m.id, m])).values()).map(
(model) => model.id,
@ -118,7 +117,13 @@ const filterModels = (
const filterFn = contextFilters[ctx.context];
if (!filterFn) return [];
return filterFn({ allProxyModels: deduplicatedProxyModels, ...extra, options: ctx.options });
const filterArgs: FilterContextArgs = {
allProxyModels: deduplicatedProxyModels,
organizationID: ctx.organizationID,
...extra,
options: ctx.options,
};
return filterFn(filterArgs);
};
export const ModelSelect = (props: ModelSelectProps) => {
@ -126,16 +131,17 @@ export const ModelSelect = (props: ModelSelectProps) => {
const { id, teamID, organizationID, options, context, dataTestId, value = [], onChange, style } = props;
const { showAllProxyModelsOverride, includeSpecialOptions } = options || {};
const { data: allProxyModels, isLoading: isLoadingAllProxyModels } = useAllProxyModels();
const { data: team, isLoading: isLoadingTeam } = useTeam(teamID);
const { data: team, isLoading: isLoadingTeam, isFetching: isFetchingTeam } = useTeam(teamID);
const { data: organization, isLoading: isLoadingOrganization } = useOrganization(organizationID);
const { data: currentUser, isLoading: isCurrentUserLoading } = useCurrentUser();
const isSpecialOption = (value: string) => MODEL_SENTINEL_OPTIONS.some((sv) => sv.value === value);
const hasSpecialOptionSelected = value.some(isSpecialOption);
const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading;
const organizationHasAllProxyModels =
organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) ||
organization?.models.length === 0;
const isTeamPending = isLoadingTeam || isAwaitingOrganizationModels(team, isFetchingTeam);
const isLoading = isLoadingAllProxyModels || isTeamPending || isLoadingOrganization || isCurrentUserLoading;
// The org's ceiling rides on /team/info, which a team admin may read; /organization/info 403s for them.
const organizationModels = team?.organization_models ?? organization?.models;
const organizationHasAllProxyModels = organizationModels !== undefined && isUncappedModelCeiling(organizationModels);
const shouldShowAllProxyModels =
showAllProxyModelsOverride || (organizationHasAllProxyModels && includeSpecialOptions) || context === "global";
@ -159,8 +165,7 @@ export const ModelSelect = (props: ModelSelectProps) => {
};
const filteredModels = filterModels(allProxyModels?.data ?? [], props, {
selectedTeam: team,
selectedOrganization: organization,
organizationModels,
userModels: currentUser?.models,
});

View file

@ -27,6 +27,8 @@ export interface Team {
access_group_models?: string[];
access_group_mcp_server_ids?: string[];
access_group_agent_ids?: string[];
// Parent org's model ceiling. undefined = no org / not loaded; [] or ["all-proxy-models"] = no ceiling.
organization_models?: string[] | null;
}
export interface KeyResponse {