mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
Fix internal tag usage scoping (#27315)
* Scope internal tag usage to own keys Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com> * Add internal tag usage unowned key regression test Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com> * Handle empty internal tag usage scopes safely Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com> * Add tag activity database guard Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com> --------- Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
ishaan-berri
GitHub
parent
5e016f9f74
commit
12e59c8798
9 changed files with 507 additions and 26 deletions
|
|
@ -677,6 +677,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",
|
||||
|
|
@ -689,7 +693,16 @@ class LiteLLMRoutes(enum.Enum):
|
|||
+ compliance_check_routes
|
||||
)
|
||||
|
||||
internal_user_view_only_routes = spend_tracking_routes + compliance_check_routes
|
||||
internal_user_view_only_routes = (
|
||||
spend_tracking_routes
|
||||
+ compliance_check_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",
|
||||
|
|
|
|||
|
|
@ -13,12 +13,12 @@ All /tag management endpoints
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import 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,
|
||||
|
|
@ -40,6 +40,72 @@ if TYPE_CHECKING:
|
|||
router = APIRouter()
|
||||
|
||||
|
||||
async def _get_internal_user_api_keys(
|
||||
prisma_client,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> List[str]:
|
||||
user_role = user_api_key_dict.user_role
|
||||
if user_role is None or not user_role.is_internal_user_role:
|
||||
return []
|
||||
|
||||
user_api_keys = set()
|
||||
if user_api_key_dict.api_key:
|
||||
user_api_keys.add(user_api_key_dict.api_key)
|
||||
|
||||
user_id = user_api_key_dict.user_id
|
||||
if user_id is None:
|
||||
return sorted(user_api_keys)
|
||||
|
||||
key_records = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"user_id": user_id},
|
||||
select={"token": True},
|
||||
)
|
||||
user_api_keys.update(
|
||||
key_record.token
|
||||
for key_record in key_records
|
||||
if getattr(key_record, "token", None)
|
||||
)
|
||||
|
||||
return sorted(user_api_keys)
|
||||
|
||||
|
||||
async def _get_tag_list_scope(
|
||||
prisma_client,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Optional[Dict[str, dict]]:
|
||||
user_role = user_api_key_dict.user_role
|
||||
if user_api_key_has_admin_view(user_api_key_dict) or (
|
||||
user_role is None or not user_role.is_internal_user_role
|
||||
):
|
||||
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]]]:
|
||||
user_role = user_api_key_dict.user_role
|
||||
if user_api_key_has_admin_view(user_api_key_dict) or (
|
||||
user_role is None or not user_role.is_internal_user_role
|
||||
):
|
||||
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:
|
||||
|
|
@ -453,9 +519,41 @@ async def list_tags(
|
|||
_validate_tag_list_date_range(start_date, end_date)
|
||||
|
||||
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: Dict[str, Any] = {"tag": {"not": None}}
|
||||
if tag_scope:
|
||||
dynamic_tag_where = {**dynamic_tag_where, **tag_scope}
|
||||
if start_date is not None and end_date is not None:
|
||||
dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date}
|
||||
|
||||
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()
|
||||
|
|
@ -489,22 +587,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_where: Dict[str, Any] = {"tag": {"not": None}}
|
||||
if start_date is not None and end_date is not None:
|
||||
dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date}
|
||||
|
||||
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},
|
||||
)
|
||||
|
||||
dynamic_tag_config = [
|
||||
{
|
||||
"name": row["tag"],
|
||||
|
|
@ -572,6 +654,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.
|
||||
|
|
@ -590,8 +673,18 @@ async def get_tag_daily_activity(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database not connected")
|
||||
|
||||
# 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,
|
||||
)
|
||||
if scoped_api_key_filter == []:
|
||||
return SpendAnalyticsPaginatedResponse(results=[])
|
||||
|
||||
return await get_daily_activity(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -602,7 +695,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
|
||||
|
|
|
|||
|
|
@ -1889,6 +1889,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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -380,6 +380,94 @@ async def test_list_tags_no_dynamic_tags():
|
|||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
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_list_tags_with_date_range_filters_dynamic_tags():
|
||||
"""
|
||||
|
|
@ -420,6 +508,209 @@ async def test_list_tags_with_date_range_filters_dynamic_tags():
|
|||
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_internal_user_tag_daily_activity_rejects_unowned_api_key_filter():
|
||||
"""
|
||||
If an internal user filters tag usage by an API key they do not own, the
|
||||
endpoint should return an empty scoped filter instead of exposing that key's
|
||||
tag spend.
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
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]
|
||||
)
|
||||
result = await get_tag_daily_activity(
|
||||
start_date="2025-01-01",
|
||||
end_date="2025-01-31",
|
||||
api_key="unowned-key",
|
||||
user_api_key_dict=mock_user_auth,
|
||||
)
|
||||
|
||||
assert result.results == []
|
||||
assert result.metadata.total_spend == 0
|
||||
assert result.metadata.total_api_requests == 0
|
||||
mock_get_daily_activity.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_tag_daily_activity_scopes_to_current_key_without_user_id():
|
||||
"""
|
||||
If an internal-user token has no user_id, it should still scope tag usage to
|
||||
the current request key instead of falling back to proxy-wide tag spend.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from litellm.proxy.management_endpoints.tag_management_endpoints import (
|
||||
get_tag_daily_activity,
|
||||
)
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
api_key="current-owned-key",
|
||||
user_id=None,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
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
|
||||
mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
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_db.litellm_verificationtoken.find_many.assert_not_awaited()
|
||||
mock_get_daily_activity.assert_awaited_once()
|
||||
assert mock_get_daily_activity.await_args.kwargs["api_key"] == [
|
||||
"current-owned-key"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_tag_daily_activity_without_any_scoped_keys_returns_empty():
|
||||
"""
|
||||
If an internal-user token has neither user_id nor api_key, the endpoint must
|
||||
return an empty response instead of dropping the API key filter.
|
||||
"""
|
||||
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=None,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
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
|
||||
mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
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.results == []
|
||||
assert result.metadata.total_spend == 0
|
||||
assert result.metadata.total_api_requests == 0
|
||||
mock_db.litellm_verificationtoken.find_many.assert_not_awaited()
|
||||
mock_get_daily_activity.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tag_daily_activity_requires_database_connection():
|
||||
"""
|
||||
Tag daily activity should fail with the same explicit DB error used by other
|
||||
tag endpoints instead of raising an AttributeError during scope resolution.
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_tag_daily_activity(
|
||||
start_date="2025-01-01",
|
||||
end_date="2025-01-31",
|
||||
user_api_key_dict=mock_user_auth,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert exc_info.value.detail == "Database not connected"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tags_without_date_range_omits_date_filter():
|
||||
"""When no date range is passed, the WHERE clause must not carry a date key."""
|
||||
|
|
|
|||
|
|
@ -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} />);
|
||||
|
||||
|
|
|
|||
|
|
@ -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("");
|
||||
|
|
@ -444,7 +445,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 && (
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue