scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

Scope internal tag usage to own keys

Browse source 
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin 2026-05-06 18:43:28 +00:00
parent b83d11351f
commit 0b76f51dc9
No known key found for this signature in database
9 changed files with 347 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/_types.py
View file

@ -666,6 +666,10 @@ class LiteLLMRoutes(enum.Enum):
"/global/activity",
"/global/activity/model",
"/global/activity/cache_hits",
# Tag usage endpoints scope internal users to tags produced by
# their own keys in tag_management_endpoints.py.
"/tag/daily/activity",
"/tag/list",
"/v1/models/{model_id}",
"/models/{model_id}",
"/guardrails/list",
@ -677,7 +681,12 @@ class LiteLLMRoutes(enum.Enum):
+ key_management_routes
)
internal_user_view_only_routes = spend_tracking_routes
internal_user_view_only_routes = spend_tracking_routes + [
# Tag usage endpoints scope internal viewers to tags produced by
# their own keys in tag_management_endpoints.py.
"/tag/daily/activity",
"/tag/list",
]
self_managed_routes = [
"/team/member_add",

130
litellm/proxy/management_endpoints/tag_management_endpoints.py
View file

@ -12,12 +12,16 @@ All /tag management endpoints
import asyncio
import json
from typing import TYPE_CHECKING, Dict, List, Optional
from typing import TYPE_CHECKING, Dict, List, Optional, Union
from fastapi import APIRouter, Depends, HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import (
LitellmUserRoles,
UserAPIKeyAuth,
user_api_key_has_admin_view,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
@ -39,6 +43,76 @@ if TYPE_CHECKING:
router = APIRouter()
def _is_internal_user_role(user_api_key_dict: UserAPIKeyAuth) -> bool:
return user_api_key_dict.user_role in (
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER.value,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
)
async def _get_internal_user_api_keys(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
) -> List[str]:
if not _is_internal_user_role(user_api_key_dict):
return []
user_id = user_api_key_dict.user_id
if user_id is None:
return []
key_records = await prisma_client.db.litellm_verificationtoken.find_many(
where={"user_id": user_id},
select={"token": True},
)
user_api_keys = {
key_record.token
for key_record in key_records
if getattr(key_record, "token", None)
}
if user_api_key_dict.api_key:
user_api_keys.add(user_api_key_dict.api_key)
return sorted(user_api_keys)
async def _get_tag_list_scope(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[Dict[str, dict]]:
if user_api_key_has_admin_view(user_api_key_dict) or not _is_internal_user_role(
user_api_key_dict
):
return None
scoped_api_keys = await _get_internal_user_api_keys(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
return {"api_key": {"in": scoped_api_keys}}
async def _get_tag_daily_activity_api_key_filter(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
requested_api_key: Optional[str],
) -> Optional[Union[str, List[str]]]:
if user_api_key_has_admin_view(user_api_key_dict) or not _is_internal_user_role(
user_api_key_dict
):
return requested_api_key
scoped_api_keys = await _get_internal_user_api_keys(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
if requested_api_key is not None:
return requested_api_key if requested_api_key in scoped_api_keys else []
return scoped_api_keys
async def _get_model_names(prisma_client, model_ids: list) -> Dict[str, str]:
"""Helper function to get model names from model IDs"""
try:
@ -412,9 +486,39 @@ async def list_tags(
raise HTTPException(status_code=500, detail="Database not connected")
try:
tag_scope = await _get_tag_list_scope(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
## QUERY DYNAMIC TAGS ##
# Use group_by instead of find_many(distinct=["tag"]).
# Prisma's distinct fetches all columns for all rows and deduplicates
# in application code, which is extremely slow on large tables.
# See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood
dynamic_tag_where = {"tag": {"not": None}}
if tag_scope:
dynamic_tag_where = {**dynamic_tag_where, **tag_scope}
dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by(
by=["tag"],
where=dynamic_tag_where,
min={"created_at": True},
max={"updated_at": True},
)
used_tag_names = [row["tag"] for row in dynamic_tag_rows if row["tag"]]
if tag_scope is not None and not used_tag_names:
return []
stored_tag_where = (
{"tag_name": {"in": used_tag_names}} if tag_scope is not None else None
)
## QUERY STORED TAGS ##
tag_records = await prisma_client.db.litellm_tagtable.find_many(
include={"litellm_budget_table": True}
where=stored_tag_where,
include={"litellm_budget_table": True},
)
stored_tag_names = set()
@ -448,18 +552,6 @@ async def list_tags(
list_of_tags.append(tag_dict)
## QUERY DYNAMIC TAGS ##
# Use group_by instead of find_many(distinct=["tag"]).
# Prisma's distinct fetches all columns for all rows and deduplicates
# in application code, which is extremely slow on large tables.
# See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood
dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by(
by=["tag"],
where={"tag": {"not": None}},
min={"created_at": True},
max={"updated_at": True},
)
dynamic_tag_config = [
{
"name": row["tag"],
@ -527,6 +619,7 @@ async def get_tag_daily_activity(
api_key: Optional[str] = None,
page: int = 1,
page_size: int = 10,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get daily activity for specific tags or all tags.
@ -547,6 +640,11 @@ async def get_tag_daily_activity(
# Convert comma-separated tags string to list if provided
tag_list = tags.split(",") if tags else None
scoped_api_key_filter = await _get_tag_daily_activity_api_key_filter(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
requested_api_key=api_key,
)
return await get_daily_activity(
prisma_client=prisma_client,
@ -557,7 +655,7 @@ async def get_tag_daily_activity(
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
api_key=scoped_api_key_filter,
page=page,
page_size=page_size,
# metadata_metrics_func=None because litellm_dailytagspend rows are

35
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -1856,6 +1856,41 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re
assert "Only proxy admin can be used to generate" in str(exc_info.value)
@pytest.mark.parametrize(
"user_role",
[
LitellmUserRoles.INTERNAL_USER.value,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
],
)
@pytest.mark.parametrize("route", ["/tag/list", "/tag/daily/activity"])
def test_internal_users_can_access_scoped_tag_usage_routes(user_role, route):
"""
Internal users can read tag usage endpoints because the endpoint handlers
scope results to the caller's own keys.
"""
user_obj = LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=user_role,
)
valid_token = UserAPIKeyAuth(
user_id="test_user",
user_role=user_role,
)
request = MagicMock(spec=Request)
request.query_params = {}
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=user_role,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
@pytest.mark.parametrize(
"user_role",
[

134
tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
View file

@ -380,6 +380,140 @@ async def test_list_tags_no_dynamic_tags():
app.dependency_overrides.clear()
@pytest.mark.asyncio
async def test_internal_user_list_tags_only_returns_tags_used_by_their_keys():
"""
Internal users can view tag usage, but the tag list must be scoped to tags
produced by API keys owned by the caller.
"""
from datetime import datetime
from unittest.mock import AsyncMock, Mock
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
mock_user_auth = UserAPIKeyAuth(
api_key="current-owned-key",
user_id="internal-user-123",
user_role=LitellmUserRoles.INTERNAL_USER,
)
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
try:
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_db = Mock()
mock_prisma.db = mock_db
owned_key_record = Mock()
owned_key_record.token = "owned-key"
mock_db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[owned_key_record]
)
mock_db.litellm_dailytagspend.group_by = AsyncMock(
return_value=[
{
"tag": "stored-owned-tag",
"_min": {"created_at": "2025-02-01T00:00:00Z"},
"_max": {"updated_at": "2025-03-01T00:00:00Z"},
},
{
"tag": "dynamic-owned-tag",
"_min": {"created_at": "2025-02-02T00:00:00Z"},
"_max": {"updated_at": "2025-03-02T00:00:00Z"},
},
]
)
stored_tag = Mock()
stored_tag.tag_name = "stored-owned-tag"
stored_tag.description = "A stored tag used by the caller"
stored_tag.models = ["model-1"]
stored_tag.model_info = {}
stored_tag.spend = 0.0
stored_tag.budget_id = None
stored_tag.created_at = datetime(2025, 1, 1)
stored_tag.updated_at = datetime(2025, 1, 1)
stored_tag.created_by = "admin-user"
stored_tag.litellm_budget_table = None
mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[stored_tag])
response = client.get(
"/tag/list",
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
assert [tag["name"] for tag in response.json()] == [
"stored-owned-tag",
"dynamic-owned-tag",
]
mock_db.litellm_verificationtoken.find_many.assert_awaited_once_with(
where={"user_id": "internal-user-123"},
select={"token": True},
)
mock_db.litellm_dailytagspend.group_by.assert_awaited_once_with(
by=["tag"],
where={
"tag": {"not": None},
"api_key": {"in": ["current-owned-key", "owned-key"]},
},
min={"created_at": True},
max={"updated_at": True},
)
mock_db.litellm_tagtable.find_many.assert_awaited_once_with(
where={"tag_name": {"in": ["stored-owned-tag", "dynamic-owned-tag"]}},
include={"litellm_budget_table": True},
)
finally:
app.dependency_overrides.clear()
@pytest.mark.asyncio
async def test_internal_user_tag_daily_activity_is_scoped_to_their_keys():
"""
Internal users must not receive proxy-wide tag spend rows when viewing tag
usage daily activity.
"""
from unittest.mock import AsyncMock, Mock
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_tag_daily_activity,
)
mock_user_auth = UserAPIKeyAuth(
user_id="internal-user-123",
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
)
with (
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch(
"litellm.proxy.management_endpoints.tag_management_endpoints.get_daily_activity",
new_callable=AsyncMock,
) as mock_get_daily_activity,
):
mock_db = Mock()
mock_prisma.db = mock_db
owned_key_record = Mock()
owned_key_record.token = "owned-key"
mock_db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[owned_key_record]
)
mock_get_daily_activity.return_value = "daily-activity-response"
result = await get_tag_daily_activity(
start_date="2025-01-01",
end_date="2025-01-31",
user_api_key_dict=mock_user_auth,
)
assert result == "daily-activity-response"
mock_get_daily_activity.assert_awaited_once()
assert mock_get_daily_activity.await_args.kwargs["api_key"] == ["owned-key"]
@pytest.mark.asyncio
async def test_get_deployments_by_model_id():
"""

30
ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx
View file

@ -57,7 +57,10 @@ vi.mock("./EndpointUsage/EndpointUsage", () => ({
vi.mock("./UsageViewSelect/UsageViewSelect", async () => {
const React = await import("react");
const UsageViewSelect = ({ value, onChange }: any) => {
const UsageViewSelect = ({ value, onChange, canViewTagUsage = false }: any) => {
const tagOption = canViewTagUsage
? React.createElement("option", { value: "tag" }, "Tag Usage")
: null;
return React.createElement(
"select",
{
@ -70,7 +73,7 @@ vi.mock("./UsageViewSelect/UsageViewSelect", async () => {
React.createElement("option", { value: "team" }, "Team Usage"),
React.createElement("option", { value: "organization" }, "Organization Usage"),
React.createElement("option", { value: "customer" }, "Customer Usage"),
React.createElement("option", { value: "tag" }, "Tag Usage"),
tagOption,
React.createElement("option", { value: "agent" }, "Agent Usage"),
React.createElement("option", { value: "user-agent-activity" }, "User Agent Activity"),
);
@ -639,6 +642,29 @@ describe("UsagePage", () => {
});
});
it("should show tag usage selector option for internal users", async () => {
mockUseAuthorized.mockReturnValue({
isLoading: false,
isAuthorized: true,
token: "mock-token",
accessToken: "test-token",
userId: "user-123",
userEmail: "test@example.com",
userRole: "internal_user",
premiumUser: true,
disabledPersonalKeyCreation: false,
showSSOBanner: false,
});
renderWithProviders(<UsagePage {...defaultProps} />);
await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});
expect(screen.getByRole("option", { name: "Tag Usage" })).toBeInTheDocument();
});
it("should show organization usage banner and view for admins", async () => {
renderWithProviders(<UsagePage {...defaultProps} organizations={mockOrganizations} />);

10
ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx
View file

@ -31,7 +31,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { all_admin_roles } from "../../../utils/roles";
import { all_admin_roles, internalUserRoles } from "../../../utils/roles";
import { ActivityMetrics, processActivityData } from "../../activity_metrics";
import CloudZeroExportModal from "../../cloudzero_export_modal";
import EntityUsageExportModal from "../../EntityUsageExport";
@ -84,6 +84,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
console.log(`currentUser: ${JSON.stringify(currentUser)}`);
console.log(`currentUser max budget: ${currentUser?.max_budget}`);
const isAdmin = all_admin_roles.includes(userRole || "");
const canViewTagUsage = isAdmin || internalUserRoles.includes(userRole || "");
// Debounced search for user selector
const [userSearchInput, setUserSearchInput] = useState("");
@ -437,7 +438,12 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<div className="flex items-end justify-between gap-6 mb-6">
<div className="flex-1">
<div className="flex items-end justify-between gap-6 mb-4 w-full">
<UsageViewSelect value={usageView} onChange={(value) => setUsageView(value)} isAdmin={isAdmin} />
<UsageViewSelect
value={usageView}
onChange={(value) => setUsageView(value)}
isAdmin={isAdmin}
canViewTagUsage={canViewTagUsage}
/>
<AdvancedDatePicker value={dateValue} onValueChange={handleDateChange} />
</div>
{paginatedResult.isFetchingMore && (

12
ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.test.tsx
View file

@ -110,4 +110,16 @@ describe("UsageViewSelect", () => {
expect(mockOnChange).toHaveBeenCalledWith("team");
});
it("should show Tag Usage for non-admin users with tag usage permission", () => {
render(<UsageViewSelect value="global" onChange={mockOnChange} isAdmin={false} canViewTagUsage={true} />);
expect(screen.getByRole("option", { name: "Tag Usage" })).toBeInTheDocument();
});
it("should hide Tag Usage for non-admin users without tag usage permission", () => {
render(<UsageViewSelect value="global" onChange={mockOnChange} isAdmin={false} />);
expect(screen.queryByRole("option", { name: "Tag Usage" })).not.toBeInTheDocument();
});
});

5
ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.tsx
View file

@ -16,6 +16,7 @@ export interface UsageViewSelectProps {
value: UsageOption;
onChange: (value: UsageOption) => void;
isAdmin: boolean;
canViewTagUsage?: boolean;
title?: string;
description?: string;
"data-id"?: string;
@ -106,12 +107,16 @@ export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
value,
onChange,
isAdmin,
canViewTagUsage = false,
title = "Usage View",
description = "Select the usage data you want to view",
"data-id": dataId,
}) => {
const getFilteredOptions = () => {
return OPTIONS.filter((option) => {
if (option.value === "tag" && canViewTagUsage) {
return true;
}
if (option.adminOnly && !isAdmin) {
return false;
}

2
ui/litellm-dashboard/src/utils/roles.ts
View file

@ -5,7 +5,7 @@ export const old_admin_roles = ["Admin", "Admin Viewer"];
export const v2_admin_role_names = ["proxy_admin", "proxy_admin_viewer", "org_admin"];
export const all_admin_roles = [...old_admin_roles, ...v2_admin_role_names];
export const internalUserRoles = ["Internal User", "Internal Viewer"];
export const internalUserRoles = ["Internal User", "Internal Viewer", "internal_user", "internal_user_viewer"];
export const rolesAllowedToSeeUsage = ["Admin", "Admin Viewer", "Internal User", "Internal Viewer"];
export const rolesWithWriteAccess = ["Internal User", "Admin", "proxy_admin"];
// Admin-tier read parity: Admin Viewer sees Models + Endpoints, Agents, and