From f62658795aeb1202df80cf290e9525de1ea85b2c Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 12:25:48 -0700 Subject: [PATCH 01/55] feat(teams): resolve access group models/MCPs/agents in team endpoints Add access_group_models, access_group_mcp_server_ids, and access_group_agent_ids to /team/info and /v2/team/list responses. These fields contain resources inherited from access groups, kept separate from direct assignments so the UI can distinguish the source. Backend: _resolve_access_group_resources() helper resolves access group resources via existing _get_*_from_access_groups() functions. UI: Teams table and detail view show direct models as blue badges and access-group-sourced models as green badges. --- litellm/proxy/_types.py | 4 + .../management_endpoints/team_endpoints.py | 54 +++++++ .../management_endpoints/team_endpoints.py | 4 + .../components/TeamsTable/ModelsCell.tsx | 142 +++++++++--------- .../components/key_team_helpers/key_list.tsx | 4 + .../src/components/team/TeamInfo.tsx | 27 +++- 6 files changed, 161 insertions(+), 74 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8faf36df4c6..6b581957e91 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3805,6 +3805,10 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): team_member_budget_table: Optional[LiteLLM_BudgetTable] = None + # Resources inherited from access groups (separate from direct assignments) + access_group_models: Optional[List[str]] = None + access_group_mcp_server_ids: Optional[List[str]] = None + access_group_agent_ids: Optional[List[str]] = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3643373be65..6b4cf72fc11 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -62,6 +62,9 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + _get_agent_ids_from_access_groups, + _get_mcp_server_ids_from_access_groups, + _get_models_from_access_groups, allowed_route_check_inside_route, can_org_access_model, get_org_object, @@ -3042,6 +3045,14 @@ async def team_info( team_info_response_object=_team_info, ) + # Resolve resources inherited from access groups + resolved = await _resolve_access_group_resources( + access_group_ids=_team_info.access_group_ids, + ) + _team_info.access_group_models = resolved["access_group_models"] + _team_info.access_group_mcp_server_ids = resolved["access_group_mcp_server_ids"] + _team_info.access_group_agent_ids = resolved["access_group_agent_ids"] + response_object = TeamInfoResponseObject( team_id=team_id, team_info=_team_info, @@ -3332,6 +3343,36 @@ async def _build_team_list_where_conditions( return where_conditions +async def _resolve_access_group_resources( + access_group_ids: Optional[List[str]], +) -> Dict[str, List[str]]: + """ + Resolve resources inherited from access groups. + + Returns only the access-group-sourced resources (not direct assignments). + Keeps them separate so callers can distinguish where each resource comes from. + """ + empty: Dict[str, List[str]] = { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + } + if not access_group_ids: + return empty + + return { + "access_group_models": await _get_models_from_access_groups( + access_group_ids=access_group_ids, + ), + "access_group_mcp_server_ids": await _get_mcp_server_ids_from_access_groups( + access_group_ids=access_group_ids, + ), + "access_group_agent_ids": await _get_agent_ids_from_access_groups( + access_group_ids=access_group_ids, + ), + } + + def _convert_teams_to_response_models( teams: list, use_deleted_table: bool, @@ -3558,6 +3599,19 @@ async def list_team_v2( # Convert Prisma models to response models with members_count team_list = _convert_teams_to_response_models(teams, use_deleted_table) + # Resolve resources inherited from access groups for each team + if not use_deleted_table: + for team_item in team_list: + if isinstance(team_item, TeamListItem): + resolved = await _resolve_access_group_resources( + access_group_ids=team_item.access_group_ids, + ) + team_item.access_group_models = resolved["access_group_models"] + team_item.access_group_mcp_server_ids = resolved[ + "access_group_mcp_server_ids" + ] + team_item.access_group_agent_ids = resolved["access_group_agent_ids"] + return { "teams": team_list, "total": total_count, diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 5055a65783f..2455eb495d1 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -47,6 +47,10 @@ class TeamListItem(LiteLLM_TeamTable): """A team item in the paginated list response, enriched with computed fields.""" members_count: int = 0 + # Resources inherited from access groups (separate from direct assignments) + access_group_models: Optional[List[str]] = None + access_group_mcp_server_ids: Optional[List[str]] = None + access_group_agent_ids: Optional[List[str]] = None class TeamListResponse(BaseModel): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx index 5cabe4c4a8f..03a0a80bc71 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx @@ -1,16 +1,54 @@ import { Badge, Icon, TableCell, Text } from "@tremor/react"; import { ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; -import React, { useState } from "react"; +import React, { useMemo, useState } from "react"; import { Team } from "@/components/key_team_helpers/key_list"; interface ModelsCellProps { team: Team; } +interface ModelEntry { + name: string; + source: "direct" | "access_group"; +} + const ModelsCell = ({ team }: ModelsCellProps) => { const [expandedAccordion, setExpandedAccordion] = useState(false); + const modelEntries: ModelEntry[] = useMemo(() => { + const entries: ModelEntry[] = (team.models || []).map((m) => ({ + name: m, + source: "direct" as const, + })); + for (const m of team.access_group_models || []) { + entries.push({ name: m, source: "access_group" }); + } + return entries; + }, [team.models, team.access_group_models]); + + const renderBadge = (entry: ModelEntry, index: number) => { + if (entry.name === "all-proxy-models") { + return ( + + All Proxy Models + + ); + } + const displayName = getModelDisplayName(entry.name); + const truncated = displayName.length > 30 ? `${displayName.slice(0, 30)}...` : displayName; + return ( + + {truncated} + + ); + }; + return ( { whiteSpace: "pre-wrap", overflow: "hidden", }} - className={team.models.length > 3 ? "px-0" : ""} + className={modelEntries.length > 3 ? "px-0" : ""} >
- {Array.isArray(team.models) ? ( + {modelEntries.length === 0 ? ( + + All Proxy Models + + ) : (
- {team.models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {team.models.length > 3 && ( -
- { - setExpandedAccordion((prev) => !prev); - }} - /> -
- )} -
- {team.models.slice(0, 3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {team.models.length > 3 && !expandedAccordion && ( - - - +{team.models.length - 3} {team.models.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordion && ( -
- {team.models.slice(3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
+
+ {modelEntries.length > 3 && ( +
+ { + setExpandedAccordion((prev) => !prev); + }} + />
- - )} + )} +
+ {modelEntries.slice(0, 3).map((entry, index) => renderBadge(entry, index))} + {modelEntries.length > 3 && !expandedAccordion && ( + + + +{modelEntries.length - 3} {modelEntries.length - 3 === 1 ? "more model" : "more models"} + + + )} + {expandedAccordion && ( +
+ {modelEntries.slice(3).map((entry, index) => renderBadge(entry, index + 3))} +
+ )} +
+
- ) : null} + )}
); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index a681e438cd1..04b9a5c9962 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -15,6 +15,10 @@ export interface Team { keys: KeyResponse[]; members_with_roles: Member[]; spend: number; + access_group_ids?: string[]; + access_group_models?: string[]; + access_group_mcp_server_ids?: string[]; + access_group_agent_ids?: string[]; } export interface KeyResponse { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index a4c7ae2bbbe..79a55cf30fd 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -655,16 +655,31 @@ const TeamInfoView: React.FC = ({ Models
- {info.models.length === 0 ? ( + {info.models.length === 0 && !(info.access_group_models?.length) ? ( All proxy models ) : ( - info.models.map((model, index) => ( - - {model} - - )) + <> + {info.models.map((model: string, index: number) => ( + + {model} + + ))} + {(info.access_group_models || []).map((model: string, index: number) => ( + + {model} + + ))} + )}
+ {info.access_group_models && info.access_group_models.length > 0 && ( +
+ + Direct + From access group + +
+ )} From bbe708b093d4fc4f59f451599eb86b991dc85b91 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 14:52:32 -0700 Subject: [PATCH 02/55] perf(teams): single-pass access group resolution + asyncio.gather in list endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fetch each access group object once and extract all 3 resource fields in a single pass instead of 3 separate calls (3N → N lookups) - Use asyncio.gather to resolve access groups across teams concurrently in list_team_v2 instead of sequential awaits - Add 5 unit tests for _resolve_access_group_resources --- .../management_endpoints/team_endpoints.py | 66 +++++-- .../test_team_endpoints.py | 163 ++++++++++++++++++ 2 files changed, 212 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6b4cf72fc11..125efb03d8c 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -62,11 +62,9 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - _get_mcp_server_ids_from_access_groups, - _get_models_from_access_groups, allowed_route_check_inside_route, can_org_access_model, + get_access_object, get_org_object, get_team_object, get_user_object, @@ -3349,6 +3347,9 @@ async def _resolve_access_group_resources( """ Resolve resources inherited from access groups. + Fetches each access group object once and extracts all three resource + fields in a single pass (models, MCP servers, agents). + Returns only the access-group-sourced resources (not direct assignments). Keeps them separate so callers can distinguish where each resource comes from. """ @@ -3360,16 +3361,38 @@ async def _resolve_access_group_resources( if not access_group_ids: return empty + from litellm.proxy.proxy_server import prisma_client as _prisma_client + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache + + if _user_api_key_cache is None: + return empty + + models: List[str] = [] + mcp_ids: List[str] = [] + agent_ids: List[str] = [] + + for ag_id in access_group_ids: + try: + ag = await get_access_object( + access_group_id=ag_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + models.extend(getattr(ag, "access_model_names", [])) + mcp_ids.extend(getattr(ag, "access_mcp_server_ids", [])) + agent_ids.extend(getattr(ag, "access_agent_ids", [])) + except Exception: + verbose_proxy_logger.debug( + "Could not fetch access group %s for resource resolution", + ag_id, + ) + return { - "access_group_models": await _get_models_from_access_groups( - access_group_ids=access_group_ids, - ), - "access_group_mcp_server_ids": await _get_mcp_server_ids_from_access_groups( - access_group_ids=access_group_ids, - ), - "access_group_agent_ids": await _get_agent_ids_from_access_groups( - access_group_ids=access_group_ids, - ), + "access_group_models": list(set(models)), + "access_group_mcp_server_ids": list(set(mcp_ids)), + "access_group_agent_ids": list(set(agent_ids)), } @@ -3601,11 +3624,20 @@ async def list_team_v2( # Resolve resources inherited from access groups for each team if not use_deleted_table: - for team_item in team_list: - if isinstance(team_item, TeamListItem): - resolved = await _resolve_access_group_resources( - access_group_ids=team_item.access_group_ids, - ) + team_items_with_ag = [ + t for t in team_list + if isinstance(t, TeamListItem) and t.access_group_ids + ] + if team_items_with_ag: + results = await asyncio.gather( + *[ + _resolve_access_group_resources( + access_group_ids=t.access_group_ids, + ) + for t in team_items_with_ag + ] + ) + for team_item, resolved in zip(team_items_with_ag, results): team_item.access_group_models = resolved["access_group_models"] team_item.access_group_mcp_server_ids = resolved[ "access_group_mcp_server_ids" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 366f659bdab..a4b5e677603 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6491,3 +6491,166 @@ async def test_create_team_member_budget_table_with_duration(): assert budget_request.budget_duration == "30d" assert budget_request.max_budget == 20.0 assert result["metadata"]["team_member_budget_id"] == "budget-abc" + + +# --------------------------------------------------------------------------- +# Tests for _resolve_access_group_resources +# --------------------------------------------------------------------------- + + +class TestResolveAccessGroupResources: + """Tests for the single-pass access group resource resolution helper.""" + + @pytest.mark.asyncio + async def test_returns_empty_when_no_access_group_ids(self): + """None or empty list should return empty lists for all resource types.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + result_none = await _resolve_access_group_resources(access_group_ids=None) + assert result_none == { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + } + + result_empty = await _resolve_access_group_resources(access_group_ids=[]) + assert result_empty == { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + } + + @pytest.mark.asyncio + async def test_single_access_group(self): + """Single access group should return its resources.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + fake_ag = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="test-group", + access_model_names=["gpt-4", "claude-3"], + access_mcp_server_ids=["mcp-1"], + access_agent_ids=["agent-1", "agent-2"], + ) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1"], + ) + + assert sorted(result["access_group_models"]) == ["claude-3", "gpt-4"] + assert result["access_group_mcp_server_ids"] == ["mcp-1"] + assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + + @pytest.mark.asyncio + async def test_multiple_access_groups_deduplicates(self): + """Multiple access groups with overlapping resources should deduplicate.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + ag1 = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="group-1", + access_model_names=["gpt-4", "claude-3"], + access_mcp_server_ids=["mcp-1"], + access_agent_ids=["agent-1"], + ) + ag2 = LiteLLM_AccessGroupTable( + access_group_id="ag-2", + access_group_name="group-2", + access_model_names=["gpt-4", "gemini"], + access_mcp_server_ids=["mcp-1", "mcp-2"], + access_agent_ids=["agent-2"], + ) + + async def fake_get_access_object(access_group_id, **kwargs): + return {"ag-1": ag1, "ag-2": ag2}[access_group_id] + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + side_effect=fake_get_access_object, + ): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1", "ag-2"], + ) + + assert sorted(result["access_group_models"]) == ["claude-3", "gemini", "gpt-4"] + assert sorted(result["access_group_mcp_server_ids"]) == ["mcp-1", "mcp-2"] + assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + + @pytest.mark.asyncio + async def test_missing_access_group_skipped(self): + """If an access group doesn't exist, it should be skipped gracefully.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + ag1 = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="group-1", + access_model_names=["gpt-4"], + access_mcp_server_ids=[], + access_agent_ids=[], + ) + + async def fake_get_access_object(access_group_id, **kwargs): + if access_group_id == "ag-1": + return ag1 + raise HTTPException(status_code=404, detail="Not found") + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + side_effect=fake_get_access_object, + ): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1", "ag-missing"], + ) + + assert result["access_group_models"] == ["gpt-4"] + assert result["access_group_mcp_server_ids"] == [] + assert result["access_group_agent_ids"] == [] + + @pytest.mark.asyncio + async def test_returns_empty_when_cache_unavailable(self): + """If user_api_key_cache is None, should return empty results.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + None, + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1"], + ) + + assert result == { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + } From 59b09102b93060b0960de0108cb5a1e457ef9188 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 15:36:19 -0700 Subject: [PATCH 03/55] docs: add default_team_params to config reference and update examples MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add default_team_params to litellm_settings reference table in config_settings.md with all sub-fields documented - Update self_serve.md and msft_sso.md examples to include team_member_permissions, tpm_limit, and rpm_limit - Fix misleading comment that implied default_team_params only applies to SSO auto-created teams — it applies to all /team/new calls --- docs/my-website/docs/proxy/config_settings.md | 1 + docs/my-website/docs/proxy/self_serve.md | 25 +++++++++++++------ docs/my-website/docs/tutorials/msft_sso.md | 10 +++++--- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index cc9090c2de6..5672e27fe51 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -201,6 +201,7 @@ router_settings: | enable_json_schema_validation | boolean | If true, enables json schema validation for all requests. | | enable_key_alias_format_validation | boolean | If true, validates `key_alias` format on `/key/generate` and `/key/update`. Must be 2-255 chars, start/end with alphanumeric, only allow `a-zA-Z0-9_-/.@`. Default `false`. | | disable_copilot_system_to_assistant | boolean | **DEPRECATED** - GitHub Copilot API supports system prompts. | +| default_team_params | object | Default parameters applied to every new team created via `/team/new` (including SSO auto-created teams). Only fills in fields not explicitly set in the request. Sub-fields: `max_budget` (float), `budget_duration` (string, e.g. `"30d"`), `models` (array of strings), `tpm_limit` (integer), `rpm_limit` (integer), `team_member_permissions` (array of strings, e.g. `["/team/daily/activity", "/key/generate"]`). | ### general_settings - Reference diff --git a/docs/my-website/docs/proxy/self_serve.md b/docs/my-website/docs/proxy/self_serve.md index b54344c1d05..7d88669bf1d 100644 --- a/docs/my-website/docs/proxy/self_serve.md +++ b/docs/my-website/docs/proxy/self_serve.md @@ -358,10 +358,15 @@ When you connect litellm to your SSO provider, litellm can auto-create teams. Us ```yaml showLineNumbers title="Default Params for new teams" litellm_settings: - default_team_params: # Default Params to apply when litellm auto creates a team from SSO IDP provider - max_budget: 100 # Optional[float], optional): $100 budget for the team - budget_duration: 30d # Optional[str], optional): 30 days budget_duration for the team - models: ["gpt-3.5-turbo"] # Optional[List[str]], optional): models to be used by the team + default_team_params: # Applied to all /team/new calls (including SSO auto-created teams) when the field is not explicitly set + max_budget: 100 # Optional[float]: $100 budget for the team + budget_duration: 30d # Optional[str]: 30 days budget_duration for the team + models: ["gpt-3.5-turbo"] # Optional[List[str]]: models to be used by the team + tpm_limit: 100000 # Optional[int]: tokens per minute limit + rpm_limit: 1000 # Optional[int]: requests per minute limit + team_member_permissions: # Optional[List[str]]: permissions granted to non-admin team members + - "/team/daily/activity" # Allow members to view team usage + - "/key/generate" # Allow members to generate API keys ``` @@ -390,10 +395,14 @@ litellm_settings: max_budget_in_team: 100 # Optional[float], optional): $100 budget for the team. Defaults to None. user_role: "user" # Optional[str], optional): "user" or "admin". Defaults to "user" - default_team_params: # Default Params to apply when litellm auto creates a team from SSO IDP provider - max_budget: 100 # Optional[float], optional): $100 budget for the team - budget_duration: 30d # Optional[str], optional): 30 days budget_duration for the team - models: ["gpt-3.5-turbo"] # Optional[List[str]], optional): models to be used by the team + default_team_params: # Applied to all /team/new calls (including SSO auto-created teams) when the field is not explicitly set + max_budget: 100 # Optional[float]: $100 budget for the team + budget_duration: 30d # Optional[str]: 30 days budget_duration for the team + models: ["gpt-3.5-turbo"] # Optional[List[str]]: models to be used by the team + tpm_limit: 100000 # Optional[int]: tokens per minute limit + rpm_limit: 1000 # Optional[int]: requests per minute limit + team_member_permissions: # Optional[List[str]]: permissions granted to non-admin team members + - "/team/daily/activity" upperbound_key_generate_params: # Upperbound for /key/generate requests when self-serve flow is on diff --git a/docs/my-website/docs/tutorials/msft_sso.md b/docs/my-website/docs/tutorials/msft_sso.md index 2936f27297f..d8b0b1b918f 100644 --- a/docs/my-website/docs/tutorials/msft_sso.md +++ b/docs/my-website/docs/tutorials/msft_sso.md @@ -123,10 +123,12 @@ Navigate to your litellm config file and set the following params ```yaml showLineNumbers title="litellm config with default_team_params" litellm_settings: - default_team_params: # Default Params to apply when litellm auto creates a team from SSO IDP provider - max_budget: 100 # Optional[float], optional): $100 budget for the team - budget_duration: 30d # Optional[str], optional): 30 days budget_duration for the team - models: ["gpt-3.5-turbo"] # Optional[List[str]], optional): models to be used by the team + default_team_params: # Applied to all /team/new calls (including SSO auto-created teams) when the field is not explicitly set + max_budget: 100 # Optional[float]: $100 budget for the team + budget_duration: 30d # Optional[str]: 30 days budget_duration for the team + models: ["gpt-3.5-turbo"] # Optional[List[str]]: models to be used by the team + team_member_permissions: # Optional[List[str]]: permissions granted to non-admin team members + - "/team/daily/activity" # Allow members to view team usage ``` ### 3.2 Auto-create a new team on LiteLLM From c19a63e2bf263f585a5e13e5073d73e20b539dc2 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 16:00:20 -0700 Subject: [PATCH 04/55] docs: clarify that models sub-field only applies to SSO auto-created teams --- docs/my-website/docs/proxy/config_settings.md | 2 +- docs/my-website/docs/proxy/self_serve.md | 4 ++-- docs/my-website/docs/tutorials/msft_sso.md | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 5672e27fe51..27d9aed52b4 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -201,7 +201,7 @@ router_settings: | enable_json_schema_validation | boolean | If true, enables json schema validation for all requests. | | enable_key_alias_format_validation | boolean | If true, validates `key_alias` format on `/key/generate` and `/key/update`. Must be 2-255 chars, start/end with alphanumeric, only allow `a-zA-Z0-9_-/.@`. Default `false`. | | disable_copilot_system_to_assistant | boolean | **DEPRECATED** - GitHub Copilot API supports system prompts. | -| default_team_params | object | Default parameters applied to every new team created via `/team/new` (including SSO auto-created teams). Only fills in fields not explicitly set in the request. Sub-fields: `max_budget` (float), `budget_duration` (string, e.g. `"30d"`), `models` (array of strings), `tpm_limit` (integer), `rpm_limit` (integer), `team_member_permissions` (array of strings, e.g. `["/team/daily/activity", "/key/generate"]`). | +| default_team_params | object | Default parameters applied to every new team created via `/team/new` (including SSO auto-created teams). Only fills in fields not explicitly set in the request. Sub-fields: `max_budget` (float), `budget_duration` (string, e.g. `"30d"`), `tpm_limit` (integer), `rpm_limit` (integer), `team_member_permissions` (array of strings, e.g. `["/team/daily/activity", "/key/generate"]`), `models` (array of strings — only applied to SSO auto-created teams). | ### general_settings - Reference diff --git a/docs/my-website/docs/proxy/self_serve.md b/docs/my-website/docs/proxy/self_serve.md index 7d88669bf1d..639cd05d019 100644 --- a/docs/my-website/docs/proxy/self_serve.md +++ b/docs/my-website/docs/proxy/self_serve.md @@ -361,7 +361,7 @@ litellm_settings: default_team_params: # Applied to all /team/new calls (including SSO auto-created teams) when the field is not explicitly set max_budget: 100 # Optional[float]: $100 budget for the team budget_duration: 30d # Optional[str]: 30 days budget_duration for the team - models: ["gpt-3.5-turbo"] # Optional[List[str]]: models to be used by the team + models: ["gpt-3.5-turbo"] # Optional[List[str]]: models for the team (only applied to SSO auto-created teams) tpm_limit: 100000 # Optional[int]: tokens per minute limit rpm_limit: 1000 # Optional[int]: requests per minute limit team_member_permissions: # Optional[List[str]]: permissions granted to non-admin team members @@ -398,7 +398,7 @@ litellm_settings: default_team_params: # Applied to all /team/new calls (including SSO auto-created teams) when the field is not explicitly set max_budget: 100 # Optional[float]: $100 budget for the team budget_duration: 30d # Optional[str]: 30 days budget_duration for the team - models: ["gpt-3.5-turbo"] # Optional[List[str]]: models to be used by the team + models: ["gpt-3.5-turbo"] # Optional[List[str]]: models for the team (only applied to SSO auto-created teams) tpm_limit: 100000 # Optional[int]: tokens per minute limit rpm_limit: 1000 # Optional[int]: requests per minute limit team_member_permissions: # Optional[List[str]]: permissions granted to non-admin team members diff --git a/docs/my-website/docs/tutorials/msft_sso.md b/docs/my-website/docs/tutorials/msft_sso.md index d8b0b1b918f..06cc2e2aa54 100644 --- a/docs/my-website/docs/tutorials/msft_sso.md +++ b/docs/my-website/docs/tutorials/msft_sso.md @@ -126,7 +126,7 @@ litellm_settings: default_team_params: # Applied to all /team/new calls (including SSO auto-created teams) when the field is not explicitly set max_budget: 100 # Optional[float]: $100 budget for the team budget_duration: 30d # Optional[str]: 30 days budget_duration for the team - models: ["gpt-3.5-turbo"] # Optional[List[str]]: models to be used by the team + models: ["gpt-3.5-turbo"] # Optional[List[str]]: models for the team (only applied to SSO auto-created teams) team_member_permissions: # Optional[List[str]]: permissions granted to non-admin team members - "/team/daily/activity" # Allow members to view team usage ``` From f0bd33486ead7ed2ac025ab3675a92969920e649 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 16:14:34 -0700 Subject: [PATCH 05/55] fix: lazy import get_access_object to break cyclic import + short-circuit all-proxy-models display - Remove get_access_object from module-level import in team_endpoints.py and use a lazy _get_access_object wrapper to avoid cyclic dependency - Add _prisma_client is None early-exit guard in _resolve_access_group_resources - Short-circuit UI to show "All Proxy Models" when team.models is empty or contains "all-proxy-models", skipping access group model resolution --- .../proxy/management_endpoints/team_endpoints.py | 16 ++++++++++++++-- .../teams/components/TeamsTable/ModelsCell.tsx | 7 +++++-- .../src/components/team/TeamInfo.tsx | 4 ++-- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 125efb03d8c..c66177aafcf 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -64,7 +64,6 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( allowed_route_check_inside_route, can_org_access_model, - get_access_object, get_org_object, get_team_object, get_user_object, @@ -111,6 +110,16 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( router = APIRouter() +def _get_access_object(*args, **kwargs): + """ + Lazily import and delegate to `get_access_object` from + `litellm.proxy.auth.auth_checks` to avoid module-level cyclic imports. + """ + from litellm.proxy.auth.auth_checks import get_access_object as _inner_get_access_object + + return _inner_get_access_object(*args, **kwargs) + + class TeamMemberBudgetHandler: """Helper class to handle team member budget, RPM, and TPM limit operations""" @@ -3368,13 +3377,16 @@ async def _resolve_access_group_resources( if _user_api_key_cache is None: return empty + if _prisma_client is None: + return empty + models: List[str] = [] mcp_ids: List[str] = [] agent_ids: List[str] = [] for ag_id in access_group_ids: try: - ag = await get_access_object( + ag = await _get_access_object( access_group_id=ag_id, prisma_client=_prisma_client, user_api_key_cache=_user_api_key_cache, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx index 03a0a80bc71..62a7fdb783f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx @@ -16,8 +16,11 @@ interface ModelEntry { const ModelsCell = ({ team }: ModelsCellProps) => { const [expandedAccordion, setExpandedAccordion] = useState(false); + const isAllModels = !team.models || team.models.length === 0 || team.models.includes("all-proxy-models"); + const modelEntries: ModelEntry[] = useMemo(() => { - const entries: ModelEntry[] = (team.models || []).map((m) => ({ + if (isAllModels) return []; + const entries: ModelEntry[] = team.models.map((m) => ({ name: m, source: "direct" as const, })); @@ -25,7 +28,7 @@ const ModelsCell = ({ team }: ModelsCellProps) => { entries.push({ name: m, source: "access_group" }); } return entries; - }, [team.models, team.access_group_models]); + }, [team.models, team.access_group_models, isAllModels]); const renderBadge = (entry: ModelEntry, index: number) => { if (entry.name === "all-proxy-models") { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 79a55cf30fd..576f0a5f99c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -655,7 +655,7 @@ const TeamInfoView: React.FC = ({ Models
- {info.models.length === 0 && !(info.access_group_models?.length) ? ( + {info.models.length === 0 || info.models.includes("all-proxy-models") ? ( All proxy models ) : ( <> @@ -672,7 +672,7 @@ const TeamInfoView: React.FC = ({ )}
- {info.access_group_models && info.access_group_models.length > 0 && ( + {info.models.length > 0 && !info.models.includes("all-proxy-models") && info.access_group_models && info.access_group_models.length > 0 && (
Direct From 7a27434f890db534bb1ebd723c925b28fb2ad275 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 17:21:54 -0700 Subject: [PATCH 06/55] feat(ui): add submit guardrail form to Submitted Guardrails tab Wire the placeholder "Add Guardrail" button to open an antd Modal+Form with team selector, guardrail name, mode, API base URL, optional extra litellm_params JSON, and optional guardrail_info JSON. Backend call is stubbed with a TODO for now. --- .../guardrails/TeamGuardrailsTab.tsx | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) diff --git a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx index a2246fd976d..c98ad18dcab 100644 --- a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx @@ -14,6 +14,7 @@ import { AlertCircleIcon, InfoIcon, } from "lucide-react"; +import { Modal, Form, Input, Select } from "antd"; import { listGuardrailSubmissions, approveGuardrailSubmission, @@ -22,6 +23,7 @@ import { type GuardrailSubmissionItem, } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import TeamDropdown from "@/components/common_components/team_dropdown"; type GuardrailStatus = "active" | "pending" | "rejected"; @@ -820,6 +822,8 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { const [isLoading, setIsLoading] = useState(true); const [error, setError] = useState(null); const [searchDebounced, setSearchDebounced] = useState(""); + const [isSubmitModalOpen, setIsSubmitModalOpen] = useState(false); + const [submitForm] = Form.useForm(); useEffect(() => { const t = setTimeout(() => setSearchDebounced(search), 300); @@ -1006,6 +1010,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) {
); } From fe8cc100de83eb8b8064eca79c0e14cd6950812f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 17:51:17 -0700 Subject: [PATCH 07/55] feat: wire submit guardrail form to POST /guardrails/register Backend: - Add required team_id to RegisterGuardrailRequest body - Validate team membership before allowing submission Frontend: - Add useRegisterGuardrail React Query mutation hook - Wire form onFinish to mutation, pass team_id from dropdown - Show notification and refresh submissions on success --- .../proxy/guardrails/guardrail_endpoints.py | 25 ++++++- .../hooks/guardrails/useRegisterGuardrail.ts | 74 +++++++++++++++++++ .../guardrails/TeamGuardrailsTab.tsx | 27 ++++--- 3 files changed, 113 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b20876ba22..2b10537d8e6 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -542,6 +542,7 @@ class RegisterGuardrailRequest(BaseModel): str, Any ] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: Optional[Dict[str, Any]] = None + team_id: str def get_litellm_params_dict(self) -> Dict[str, Any]: return dict(self.litellm_params) @@ -603,11 +604,29 @@ async def register_guardrail( if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") - if not user_api_key_dict.team_id: + if not request.team_id: raise HTTPException( status_code=400, - detail="Registration requires an API key associated with a team. Use a team-scoped key.", + detail="team_id is required.", ) + team_id = request.team_id + + # Validate the user is a member of the specified team + if team_id != user_api_key_dict.team_id: + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.proxy_server import user_api_key_cache + + membership = await get_team_membership( + user_id=user_api_key_dict.user_id or "", + team_id=request.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if membership is None: + raise HTTPException( + status_code=403, + detail=f"You are not a member of team {request.team_id!r}", + ) params = request.get_litellm_params_dict() if params.get("guardrail") != GENERIC_GUARDRAIL_API: @@ -673,7 +692,7 @@ async def register_guardrail( "litellm_params": litellm_params_str, "guardrail_info": guardrail_info_str, "status": "pending_review", - "team_id": user_api_key_dict.team_id, + "team_id": team_id, "submitted_at": now, "created_at": now, "updated_at": now, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts new file mode 100644 index 00000000000..ccc26c4b17f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts @@ -0,0 +1,74 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface RegisterGuardrailParams { + guardrail_name: string; + litellm_params: Record; + guardrail_info?: Record; + team_id: string; +} + +export interface RegisterGuardrailResponse { + guardrail_id: string; + guardrail_name: string; + status: string; + submitted_at?: string | null; +} + +// ── Fetch function ─────────────────────────────────────────────────────────── + +const registerGuardrail = async ( + accessToken: string, + params: RegisterGuardrailParams, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl}/guardrails/register`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(params), + }); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ───────────────────────────────────────────────────────────────────── + +const guardrailKeys = createQueryKeys("guardrails"); + +export const useRegisterGuardrail = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (params) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return registerGuardrail(accessToken, params); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: guardrailKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx index c98ad18dcab..fd3f8b15618 100644 --- a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx @@ -24,6 +24,7 @@ import { } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; import TeamDropdown from "@/components/common_components/team_dropdown"; +import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail"; type GuardrailStatus = "active" | "pending" | "rejected"; @@ -824,6 +825,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { const [searchDebounced, setSearchDebounced] = useState(""); const [isSubmitModalOpen, setIsSubmitModalOpen] = useState(false); const [submitForm] = Form.useForm(); + const registerGuardrail = useRegisterGuardrail(); useEffect(() => { const t = setTimeout(() => setSearchDebounced(search), 300); @@ -1099,22 +1101,27 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { form={submitForm} layout="vertical" initialValues={{ mode: "pre_call" }} - onFinish={(values) => { + onFinish={async (values) => { const litellm_params: Record = { guardrail: "generic_guardrail_api", mode: values.mode, api_base: values.api_base, ...(values.extra_litellm_params ? JSON.parse(values.extra_litellm_params) : {}), }; - const payload = { - guardrail_name: values.guardrail_name, - litellm_params, - guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined, - }; - // TODO: call registerGuardrailCall once backend is wired - console.log("Submit guardrail:", payload); - setIsSubmitModalOpen(false); - submitForm.resetFields(); + try { + await registerGuardrail.mutateAsync({ + team_id: values.team_id, + guardrail_name: values.guardrail_name, + litellm_params, + guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined, + }); + NotificationsManager.success("Guardrail submitted for review"); + setIsSubmitModalOpen(false); + submitForm.resetFields(); + fetchSubmissions(); + } catch { + // error already handled by networking layer + } }} > Date: Thu, 2 Apr 2026 18:00:58 -0700 Subject: [PATCH 08/55] feat: wire guardrail submission to backend and allow internal user access Backend: - Add required team_id to RegisterGuardrailRequest body - Validate team membership before allowing submission Frontend: - Add useRegisterGuardrail React Query mutation hook - Wire form onFinish to mutation, pass team_id from dropdown - Remove admin-only restriction on guardrails nav item - Internal users see only Test Playground + Submitted Guardrails tabs - Admins continue to see all 4 tabs --- .../src/components/guardrails.tsx | 174 +++++++++--------- .../src/components/leftnav.tsx | 1 - 2 files changed, 88 insertions(+), 87 deletions(-) diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 56bba2724d0..6bb4b387a88 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -137,103 +137,105 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole
- Guardrail Garden - Guardrails + {isAdmin && Guardrail Garden} + {isAdmin && Guardrails} Test Playground Submitted Guardrails - {/* Guardrail Garden Tab */} - - - - - {/* Existing Guardrails Tab */} - -
- , - label: "Add Provider Guardrail", - onClick: handleAddGuardrail, - }, - { - key: "custom_code", - icon: , - label: "Create Custom Code Guardrail", - onClick: handleAddCustomCodeGuardrail, - }, - ], - }} - trigger={["click"]} - disabled={!accessToken} - > - - -
- - {selectedGuardrailId ? ( - setSelectedGuardrailId(null)} + {isAdmin && ( + + - ) : ( - + )} + + {isAdmin && ( + +
+ , + label: "Add Provider Guardrail", + onClick: handleAddGuardrail, + }, + { + key: "custom_code", + icon: , + label: "Create Custom Code Guardrail", + onClick: handleAddCustomCodeGuardrail, + }, + ], + }} + trigger={["click"]} + disabled={!accessToken} + > + + +
+ + {selectedGuardrailId ? ( + setSelectedGuardrailId(null)} + accessToken={accessToken} + isAdmin={isAdmin} + /> + ) : ( + setSelectedGuardrailId(id)} + /> + )} + + setSelectedGuardrailId(id)} + onSuccess={handleSuccess} /> - )} - + - - - - + + + )} {/* Test Playground Tab */} diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 48df5cdde80..9005311052e 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -137,7 +137,6 @@ const menuGroups: MenuGroup[] = [ page: "guardrails", label: "Guardrails", icon: , - roles: all_admin_roles, }, { key: "policies", From adb7454f857f39d22442c422a6436f29a912817b Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 18:08:22 -0700 Subject: [PATCH 09/55] fix: address review feedback on guardrail registration - Make team_id optional again, fall back to API key's team_id (preserves backwards compatibility) - Add PROXY_ADMIN bypass for team membership check (admins can register guardrails for any team) - Move get_team_membership import to module level - Prevent extra_litellm_params spread from overwriting controlled fields (guardrail, mode, api_base) by spreading extras first --- .../proxy/guardrails/guardrail_endpoints.py | 20 ++++++++++--------- .../hooks/guardrails/useRegisterGuardrail.ts | 2 +- .../guardrails/TeamGuardrailsTab.tsx | 2 +- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b10537d8e6..b1322d0eefd 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -17,6 +17,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import ( CustomCodeValidationError, @@ -542,7 +543,7 @@ class RegisterGuardrailRequest(BaseModel): str, Any ] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: Optional[Dict[str, Any]] = None - team_id: str + team_id: Optional[str] = None def get_litellm_params_dict(self) -> Dict[str, Any]: return dict(self.litellm_params) @@ -604,28 +605,29 @@ async def register_guardrail( if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") - if not request.team_id: + # Resolve team_id: prefer request body, fall back to API key's team + team_id = request.team_id or user_api_key_dict.team_id + if not team_id: raise HTTPException( status_code=400, - detail="team_id is required.", + detail="team_id is required. Provide it in the request body or use a team-scoped API key.", ) - team_id = request.team_id - # Validate the user is a member of the specified team - if team_id != user_api_key_dict.team_id: - from litellm.proxy.auth.auth_checks import get_team_membership + # Validate team membership for non-admin users when team differs from key + is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + if not is_admin and team_id != user_api_key_dict.team_id: from litellm.proxy.proxy_server import user_api_key_cache membership = await get_team_membership( user_id=user_api_key_dict.user_id or "", - team_id=request.team_id, + team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, ) if membership is None: raise HTTPException( status_code=403, - detail=f"You are not a member of team {request.team_id!r}", + detail=f"You are not a member of team {team_id!r}", ) params = request.get_litellm_params_dict() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts index ccc26c4b17f..3135e8326fc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts @@ -14,7 +14,7 @@ export interface RegisterGuardrailParams { guardrail_name: string; litellm_params: Record; guardrail_info?: Record; - team_id: string; + team_id?: string; } export interface RegisterGuardrailResponse { diff --git a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx index fd3f8b15618..8bd2cdd1964 100644 --- a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx @@ -1103,10 +1103,10 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { initialValues={{ mode: "pre_call" }} onFinish={async (values) => { const litellm_params: Record = { + ...(values.extra_litellm_params ? JSON.parse(values.extra_litellm_params) : {}), guardrail: "generic_guardrail_api", mode: values.mode, api_base: values.api_base, - ...(values.extra_litellm_params ? JSON.parse(values.extra_litellm_params) : {}), }; try { await registerGuardrail.mutateAsync({ From a287154905df6812553a983215dbbfd18f489e21 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 09:34:29 -0700 Subject: [PATCH 10/55] fix(ui): migrate guardrails tabs to antd and fix internal user view - Replace Tremor TabGroup with antd Tabs (key-based matching) to fix blank content when conditional tabs are hidden for non-admins - Skip /guardrails/submissions fetch for non-admin users via useAuthorized() hook (avoids 401 error) - Move get_team_membership to inline import in register endpoint - Non-admins default to Submitted Guardrails tab with Add Guardrail button visible --- .../proxy/guardrails/guardrail_endpoints.py | 2 +- .../src/components/guardrails.tsx | 242 +++++++++--------- .../guardrails/TeamGuardrailsTab.tsx | 8 +- 3 files changed, 130 insertions(+), 122 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index b1322d0eefd..b88c6524bb7 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -17,7 +17,6 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import ( CustomCodeValidationError, @@ -616,6 +615,7 @@ async def register_guardrail( # Validate team membership for non-admin users when team differs from key is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN if not is_admin and team_id != user_api_key_dict.team_id: + from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.proxy_server import user_api_key_cache membership = await get_team_membership( diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 6bb4b387a88..d52aba15ab0 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -1,6 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; -import { Dropdown } from "antd"; +import { Button } from "@tremor/react"; +import { Dropdown, Tabs } from "antd"; import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons"; import { getGuardrailsList, deleteGuardrailCall } from "./networking"; import AddGuardrailForm from "./guardrails/add_guardrail_form"; @@ -48,8 +48,6 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [guardrailToDelete, setGuardrailToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [selectedGuardrailId, setSelectedGuardrailId] = useState(null); - const [activeTab, setActiveTab] = useState(0); - const isAdmin = userRole ? isAdminRole(userRole) : false; const fetchGuardrails = async () => { @@ -135,124 +133,130 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole return (
- - - {isAdmin && Guardrail Garden} - {isAdmin && Guardrails} - Test Playground - Submitted Guardrails - + + ), + }, + { + key: "guardrails", + label: "Guardrails", + children: ( + <> +
+ , + label: "Add Provider Guardrail", + onClick: handleAddGuardrail, + }, + { + key: "custom_code", + icon: , + label: "Create Custom Code Guardrail", + onClick: handleAddCustomCodeGuardrail, + }, + ], + }} + trigger={["click"]} + disabled={!accessToken} + > + + +
- - {isAdmin && ( - - setSelectedGuardrailId(null)} + accessToken={accessToken} + isAdmin={isAdmin} + /> + ) : ( + setSelectedGuardrailId(id)} + /> + )} + + + + + + + + ), + }, + ] + : []), + { + key: "playground", + label: "Test Playground", + disabled: !accessToken || guardrailsList.length === 0, + children: ( + {}} /> - - )} - - {isAdmin && ( - -
- , - label: "Add Provider Guardrail", - onClick: handleAddGuardrail, - }, - { - key: "custom_code", - icon: , - label: "Create Custom Code Guardrail", - onClick: handleAddCustomCodeGuardrail, - }, - ], - }} - trigger={["click"]} - disabled={!accessToken} - > - - -
- - {selectedGuardrailId ? ( - setSelectedGuardrailId(null)} - accessToken={accessToken} - isAdmin={isAdmin} - /> - ) : ( - setSelectedGuardrailId(id)} - /> - )} - - - - - - - - )} - - {/* Test Playground Tab */} - - setActiveTab(0)} - /> - - - {/* Team Guardrails Tab */} - - - - - + ), + }, + { + key: "submitted", + label: "Submitted Guardrails", + children: , + }, + ]} + />
); }; diff --git a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx index 8bd2cdd1964..c3c1836e3c6 100644 --- a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx @@ -25,6 +25,8 @@ import { import NotificationsManager from "@/components/molecules/notifications_manager"; import TeamDropdown from "@/components/common_components/team_dropdown"; import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { isAdminRole } from "@/utils/roles"; type GuardrailStatus = "active" | "pending" | "rejected"; @@ -803,6 +805,8 @@ interface TeamGuardrailsTabProps { } export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { + const { userRole } = useAuthorized(); + const isAdmin = userRole ? isAdminRole(userRole) : false; const [guardrails, setGuardrails] = useState([]); const [summary, setSummary] = useState({ total: 0, @@ -833,7 +837,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { }, [search]); const fetchSubmissions = useCallback(async () => { - if (!accessToken) { + if (!accessToken || !isAdmin) { setIsLoading(false); return; } @@ -858,7 +862,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { } finally { setIsLoading(false); } - }, [accessToken, statusFilter, searchDebounced]); + }, [accessToken, isAdmin, statusFilter, searchDebounced]); useEffect(() => { fetchSubmissions(); From 5a8f910fe3c35d14ff1e20c0429e652faca227ab Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 11:54:52 -0700 Subject: [PATCH 11/55] add: making organizations a select instead of read only badges --- .../src/components/team/TeamInfo.tsx | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index a4c7ae2bbbe..688adb56f4a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -491,7 +491,7 @@ const TeamInfoView: React.FC = ({ ...(secretManagerSettings !== undefined ? { secret_manager_settings: secretManagerSettings } : {}), }, ...(values.policies?.length > 0 ? { policies: values.policies } : {}), - organization_id: values.organization_id, + organization_id: values.organization_id ?? "", }; updateData.max_budget = mapEmptyStringToNull(updateData.max_budget); @@ -1077,8 +1077,17 @@ const TeamInfoView: React.FC = ({ /> - - + + onChange?.(val)} + disabled={disabled} + allowClear + filterOption={false} + onSearch={handleSearch} + searchValue={searchInput} + onPopupScroll={handlePopupScroll} + loading={isLoading} + notFoundContent={isLoading ? : "No teams found"} + style={{ width: "100%" }} + popupRender={(menu) => ( + <> + {menu} + {isFetchingNextPage && ( +
+ +
+ )} + + )} + > + {teams.map((team) => ( + + {team.team_alias}{" "} + ({team.team_id}) + + ))} + + ); +}; + +export default TeamMultiSelect; From 1533f6896e51c9848b2ef9b4e95f194f89a2a26f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 14:36:46 -0700 Subject: [PATCH 14/55] fix(ui): fix imports and update placeholder for team multi select --- .../src/components/common_components/team_multi_select.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx index eec08c58dbb..6f00a4d1f7b 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx @@ -25,7 +25,7 @@ const TeamMultiSelect: React.FC = ({ disabled, organizationId, pageSize = 20, - placeholder = "Search teams by name or ID...", + placeholder = "Search teams by alias...", }) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { From 96b660b25766a9b390d962728f1c2db389663e29 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 15:38:16 -0700 Subject: [PATCH 15/55] fix(ui): wire team_id filter to key alias dropdown on Virtual Keys tab The Key Alias dropdown on the Virtual Keys page was showing aliases from all teams regardless of which team was selected. The team_id was never passed through the frontend chain to the backend /key/aliases endpoint. - Backend: add optional team_id query param to /key/aliases endpoint - networking.tsx: add team_id param to keyAliasesCall - useKeyAliases: accept and forward team_id to API call and query key - filter.tsx: pass allFilters context to custom filter components - PaginatedKeyAliasSelect: read Team ID from allFilters and pass to hook --- .../proxy/management_endpoints/key_management_endpoints.py | 7 +++++++ .../src/app/(dashboard)/hooks/keys/useKeyAliases.ts | 3 +++ .../PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx | 6 +++++- ui/litellm-dashboard/src/components/molecules/filter.tsx | 2 ++ ui/litellm-dashboard/src/components/networking.tsx | 2 ++ 5 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 831922ec3f9..fa521ac55f6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4345,6 +4345,9 @@ async def key_aliases( search: Optional[str] = Query( None, description="Search key aliases (case-insensitive partial match)" ), + team_id: Optional[str] = Query( + None, description="Filter aliases to keys belonging to this team" + ), ) -> Dict[str, Any]: """ Lists key aliases with pagination and optional search. @@ -4420,6 +4423,10 @@ async def key_aliases( query_params.append(f"%{search}%") where_parts.append(f"key_alias ILIKE ${len(query_params)}") + if team_id: + query_params.append(team_id) + where_parts.append(f"team_id = ${len(query_params)}") + where_sql = " AND ".join(where_parts) count_sql = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts index f67b15f3a9f..03e96fe73c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts @@ -8,6 +8,7 @@ const infiniteKeyAliasKeys = createQueryKeys("infiniteKeyAliases"); export const useInfiniteKeyAliases = ( size: number = 50, search?: string, + team_id?: string, ) => { const { accessToken } = useAuthorized(); return useInfiniteQuery({ @@ -15,6 +16,7 @@ export const useInfiniteKeyAliases = ( filters: { size, ...(search && { search }), + ...(team_id && { team_id }), }, }), queryFn: async ({ pageParam }) => { @@ -23,6 +25,7 @@ export const useInfiniteKeyAliases = ( pageParam as number, size, search, + team_id, ); }, initialPageParam: 1, diff --git a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx index 0bec77ca52b..940f0b7e951 100644 --- a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx +++ b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx @@ -12,6 +12,7 @@ export interface PaginatedKeyAliasSelectProps { pageSize?: number; allowClear?: boolean; disabled?: boolean; + allFilters?: { [key: string]: string }; } const SCROLL_THRESHOLD = 0.8; @@ -25,19 +26,22 @@ export const PaginatedKeyAliasSelect = ({ pageSize = 50, allowClear = true, disabled = false, + allFilters, }: PaginatedKeyAliasSelectProps) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { wait: DEBOUNCE_MS, }); + const teamId = allFilters?.["Team ID"] || undefined; + const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading, - } = useInfiniteKeyAliases(pageSize, debouncedSearch || undefined); + } = useInfiniteKeyAliases(pageSize, debouncedSearch || undefined, teamId); const options = useMemo(() => { if (!data?.pages) return []; diff --git a/ui/litellm-dashboard/src/components/molecules/filter.tsx b/ui/litellm-dashboard/src/components/molecules/filter.tsx index dcf22293f86..34ff1983f36 100644 --- a/ui/litellm-dashboard/src/components/molecules/filter.tsx +++ b/ui/litellm-dashboard/src/components/molecules/filter.tsx @@ -7,6 +7,7 @@ export interface FilterOptionCustomComponentProps { value?: string; onChange: (value: string) => void; placeholder?: string; + allFilters?: { [key: string]: string }; } export interface FilterOption { @@ -209,6 +210,7 @@ const FilterComponent: React.FC = ({ value={tempValues[option.name] || undefined} onChange={(value) => handleFilterChange(option.name, value ?? "")} placeholder={`Select ${option.label || option.name}...`} + allFilters={tempValues} /> ); })() diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 2e8518d00a6..33860e991db 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3263,6 +3263,7 @@ export const keyAliasesCall = async ( page: number = 1, size: number = 50, search?: string, + team_id?: string, ): Promise => { /** * Get key aliases from proxy with pagination and optional search @@ -3273,6 +3274,7 @@ export const keyAliasesCall = async ( page: String(page), size: String(size), ...(search ? { search } : {}), + ...(team_id ? { team_id } : {}), }), ); let url = proxyBaseUrl ? `${proxyBaseUrl}/key/aliases` : `/key/aliases`; From 38f6c9491d602dd3b39938d1a5543806be2d2735 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 16:16:55 -0700 Subject: [PATCH 16/55] fix(tests): correct mock targets in TestResolveAccessGroupResources Three tests were patching the non-existent `get_access_object` instead of `_get_access_object` (the lazy-import wrapper), causing AttributeError. Also added missing `prisma_client` mock so tests get past the early-exit guard and actually exercise the resolution logic. --- .../test_team_endpoints.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a4b5e677603..a2c8c4d560c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6539,7 +6539,7 @@ class TestResolveAccessGroupResources: ) with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + "litellm.proxy.management_endpoints.team_endpoints._get_access_object", new_callable=AsyncMock, return_value=fake_ag, ): @@ -6547,9 +6547,13 @@ class TestResolveAccessGroupResources: "litellm.proxy.proxy_server.user_api_key_cache", MagicMock(), ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1"], - ) + with patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1"], + ) assert sorted(result["access_group_models"]) == ["claude-3", "gpt-4"] assert result["access_group_mcp_server_ids"] == ["mcp-1"] @@ -6582,16 +6586,20 @@ class TestResolveAccessGroupResources: return {"ag-1": ag1, "ag-2": ag2}[access_group_id] with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + "litellm.proxy.management_endpoints.team_endpoints._get_access_object", side_effect=fake_get_access_object, ): with patch( "litellm.proxy.proxy_server.user_api_key_cache", MagicMock(), ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1", "ag-2"], - ) + with patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1", "ag-2"], + ) assert sorted(result["access_group_models"]) == ["claude-3", "gemini", "gpt-4"] assert sorted(result["access_group_mcp_server_ids"]) == ["mcp-1", "mcp-2"] @@ -6619,16 +6627,20 @@ class TestResolveAccessGroupResources: raise HTTPException(status_code=404, detail="Not found") with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + "litellm.proxy.management_endpoints.team_endpoints._get_access_object", side_effect=fake_get_access_object, ): with patch( "litellm.proxy.proxy_server.user_api_key_cache", MagicMock(), ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1", "ag-missing"], - ) + with patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1", "ag-missing"], + ) assert result["access_group_models"] == ["gpt-4"] assert result["access_group_mcp_server_ids"] == [] From ea32cb58a8a9cc12eeef9f6e5db6ff1f0666b112 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 16:20:54 -0700 Subject: [PATCH 17/55] fix: use direct attribute access with or [] fallback in _resolve_access_group_resources Replace getattr(ag, "field", []) with ag.field or [] for cleaner access and safe handling if a field is None. --- litellm/proxy/management_endpoints/team_endpoints.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c66177aafcf..ac606a71ba6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3392,9 +3392,9 @@ async def _resolve_access_group_resources( user_api_key_cache=_user_api_key_cache, proxy_logging_obj=_proxy_logging_obj, ) - models.extend(getattr(ag, "access_model_names", [])) - mcp_ids.extend(getattr(ag, "access_mcp_server_ids", [])) - agent_ids.extend(getattr(ag, "access_agent_ids", [])) + models.extend(ag.access_model_names or []) + mcp_ids.extend(ag.access_mcp_server_ids or []) + agent_ids.extend(ag.access_agent_ids or []) except Exception: verbose_proxy_logger.debug( "Could not fetch access group %s for resource resolution", From 3bdd04250784a10ce031c35c0cd50e7f03001243 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 16:49:55 -0700 Subject: [PATCH 18/55] fix(ui): remove model source legend from team detail view The blue/green color distinction is self-explanatory; the legend added visual clutter without providing enough value. --- ui/litellm-dashboard/src/components/team/TeamInfo.tsx | 8 -------- 1 file changed, 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 576f0a5f99c..4bb1311ebf6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -672,14 +672,6 @@ const TeamInfoView: React.FC = ({ )}
- {info.models.length > 0 && !info.models.includes("all-proxy-models") && info.access_group_models && info.access_group_models.length > 0 && ( -
- - Direct - From access group - -
- )} From bb03a11d7c43dd36c6a2b075afd780670b11a42e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 16:59:47 -0700 Subject: [PATCH 19/55] fix(ui): add missing access_group fields to TeamData.team_info type The TeamData interface was missing access_group_models, access_group_mcp_server_ids, and access_group_agent_ids fields, causing a TypeScript build failure. --- ui/litellm-dashboard/src/components/team/TeamInfo.tsx | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 4bb1311ebf6..1e4345905ed 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -95,6 +95,9 @@ export interface TeamData { } | null; created_at: string; access_group_ids?: string[]; + access_group_models?: string[]; + access_group_mcp_server_ids?: string[]; + access_group_agent_ids?: string[]; guardrails?: string[]; policies?: string[]; object_permission?: { From 93369bf60d7ad87f33c56ca66ce20a3994419e8d Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 3 Apr 2026 17:13:56 -0700 Subject: [PATCH 20/55] perf(teams): batch-fetch access groups in single DB query Replace per-ID _resolve_access_group_resources loop with a single find_many call that deduplicates IDs across all teams. Removes the N+1 query pattern on cold cache for the team list endpoint. --- .../management_endpoints/team_endpoints.py | 129 ++++------- .../test_team_endpoints.py | 219 +++++++----------- 2 files changed, 135 insertions(+), 213 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ac606a71ba6..ae89f09bb44 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -110,16 +110,6 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( router = APIRouter() -def _get_access_object(*args, **kwargs): - """ - Lazily import and delegate to `get_access_object` from - `litellm.proxy.auth.auth_checks` to avoid module-level cyclic imports. - """ - from litellm.proxy.auth.auth_checks import get_access_object as _inner_get_access_object - - return _inner_get_access_object(*args, **kwargs) - - class TeamMemberBudgetHandler: """Helper class to handle team member budget, RPM, and TPM limit operations""" @@ -3053,12 +3043,17 @@ async def team_info( ) # Resolve resources inherited from access groups - resolved = await _resolve_access_group_resources( - access_group_ids=_team_info.access_group_ids, - ) - _team_info.access_group_models = resolved["access_group_models"] - _team_info.access_group_mcp_server_ids = resolved["access_group_mcp_server_ids"] - _team_info.access_group_agent_ids = resolved["access_group_agent_ids"] + if _team_info.access_group_ids: + ag_lookup = await _batch_resolve_access_group_resources(_team_info.access_group_ids) + models, mcp_ids, agent_ids = set(), set(), set() + for ag_id in _team_info.access_group_ids: + if ag_id in ag_lookup: + models.update(ag_lookup[ag_id]["models"]) + mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) + agent_ids.update(ag_lookup[ag_id]["agent_ids"]) + _team_info.access_group_models = list(models) + _team_info.access_group_mcp_server_ids = list(mcp_ids) + _team_info.access_group_agent_ids = list(agent_ids) response_object = TeamInfoResponseObject( team_id=team_id, @@ -3350,62 +3345,34 @@ async def _build_team_list_where_conditions( return where_conditions -async def _resolve_access_group_resources( - access_group_ids: Optional[List[str]], -) -> Dict[str, List[str]]: +async def _batch_resolve_access_group_resources( + all_access_group_ids: List[str], +) -> Dict[str, Dict[str, List[str]]]: """ - Resolve resources inherited from access groups. + Batch-fetch access groups in a single DB query and return a per-group + resource map. - Fetches each access group object once and extracts all three resource - fields in a single pass (models, MCP servers, agents). - - Returns only the access-group-sourced resources (not direct assignments). - Keeps them separate so callers can distinguish where each resource comes from. + Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}. + Missing/invalid groups are silently omitted. """ - empty: Dict[str, List[str]] = { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } - if not access_group_ids: - return empty - from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache - if _user_api_key_cache is None: - return empty + if not all_access_group_ids or _prisma_client is None: + return {} - if _prisma_client is None: - return empty + unique_ids = list(set(all_access_group_ids)) + rows = await _prisma_client.db.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": unique_ids}}, + ) - models: List[str] = [] - mcp_ids: List[str] = [] - agent_ids: List[str] = [] - - for ag_id in access_group_ids: - try: - ag = await _get_access_object( - access_group_id=ag_id, - prisma_client=_prisma_client, - user_api_key_cache=_user_api_key_cache, - proxy_logging_obj=_proxy_logging_obj, - ) - models.extend(ag.access_model_names or []) - mcp_ids.extend(ag.access_mcp_server_ids or []) - agent_ids.extend(ag.access_agent_ids or []) - except Exception: - verbose_proxy_logger.debug( - "Could not fetch access group %s for resource resolution", - ag_id, - ) - - return { - "access_group_models": list(set(models)), - "access_group_mcp_server_ids": list(set(mcp_ids)), - "access_group_agent_ids": list(set(agent_ids)), - } + result: Dict[str, Dict[str, List[str]]] = {} + for row in rows: + result[row.access_group_id] = { + "models": list(row.access_model_names or []), + "mcp_server_ids": list(row.access_mcp_server_ids or []), + "agent_ids": list(row.access_agent_ids or []), + } + return result def _convert_teams_to_response_models( @@ -3634,27 +3601,29 @@ async def list_team_v2( # Convert Prisma models to response models with members_count team_list = _convert_teams_to_response_models(teams, use_deleted_table) - # Resolve resources inherited from access groups for each team + # Resolve resources inherited from access groups (single batch query) if not use_deleted_table: team_items_with_ag = [ t for t in team_list if isinstance(t, TeamListItem) and t.access_group_ids ] if team_items_with_ag: - results = await asyncio.gather( - *[ - _resolve_access_group_resources( - access_group_ids=t.access_group_ids, - ) - for t in team_items_with_ag - ] - ) - for team_item, resolved in zip(team_items_with_ag, results): - team_item.access_group_models = resolved["access_group_models"] - team_item.access_group_mcp_server_ids = resolved[ - "access_group_mcp_server_ids" - ] - team_item.access_group_agent_ids = resolved["access_group_agent_ids"] + all_ag_ids = [ + ag_id + for t in team_items_with_ag + for ag_id in (t.access_group_ids or []) + ] + ag_lookup = await _batch_resolve_access_group_resources(all_ag_ids) + for team_item in team_items_with_ag: + models, mcp_ids, agent_ids = set(), set(), set() + for ag_id in (team_item.access_group_ids or []): + if ag_id in ag_lookup: + models.update(ag_lookup[ag_id]["models"]) + mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) + agent_ids.update(ag_lookup[ag_id]["agent_ids"]) + team_item.access_group_models = list(models) + team_item.access_group_mcp_server_ids = list(mcp_ids) + team_item.access_group_agent_ids = list(agent_ids) return { "teams": team_list, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a2c8c4d560c..8f0a045cf86 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6494,175 +6494,128 @@ async def test_create_team_member_budget_table_with_duration(): # --------------------------------------------------------------------------- -# Tests for _resolve_access_group_resources +# Tests for _batch_resolve_access_group_resources # --------------------------------------------------------------------------- -class TestResolveAccessGroupResources: - """Tests for the single-pass access group resource resolution helper.""" +class TestBatchResolveAccessGroupResources: + """Tests for the batch access group resource resolution helper.""" @pytest.mark.asyncio - async def test_returns_empty_when_no_access_group_ids(self): - """None or empty list should return empty lists for all resource types.""" + async def test_returns_empty_when_no_ids(self): + """Empty list should return empty dict.""" from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - result_none = await _resolve_access_group_resources(access_group_ids=None) - assert result_none == { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } - - result_empty = await _resolve_access_group_resources(access_group_ids=[]) - assert result_empty == { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } + assert await _batch_resolve_access_group_resources([]) == {} @pytest.mark.asyncio async def test_single_access_group(self): """Single access group should return its resources.""" - from litellm.proxy._types import LiteLLM_AccessGroupTable from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - fake_ag = LiteLLM_AccessGroupTable( - access_group_id="ag-1", - access_group_name="test-group", - access_model_names=["gpt-4", "claude-3"], - access_mcp_server_ids=["mcp-1"], - access_agent_ids=["agent-1", "agent-2"], - ) + fake_row = MagicMock() + fake_row.access_group_id = "ag-1" + fake_row.access_model_names = ["gpt-4", "claude-3"] + fake_row.access_mcp_server_ids = ["mcp-1"] + fake_row.access_agent_ids = ["agent-1", "agent-2"] - with patch( - "litellm.proxy.management_endpoints.team_endpoints._get_access_object", - new_callable=AsyncMock, - return_value=fake_ag, - ): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - MagicMock(), - ): - with patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1"], - ) + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[fake_row]) - assert sorted(result["access_group_models"]) == ["claude-3", "gpt-4"] - assert result["access_group_mcp_server_ids"] == ["mcp-1"] - assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1"]) + + assert sorted(result["ag-1"]["models"]) == ["claude-3", "gpt-4"] + assert result["ag-1"]["mcp_server_ids"] == ["mcp-1"] + assert sorted(result["ag-1"]["agent_ids"]) == ["agent-1", "agent-2"] @pytest.mark.asyncio - async def test_multiple_access_groups_deduplicates(self): - """Multiple access groups with overlapping resources should deduplicate.""" - from litellm.proxy._types import LiteLLM_AccessGroupTable + async def test_multiple_access_groups(self): + """Multiple access groups returned in a single query.""" from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - ag1 = LiteLLM_AccessGroupTable( - access_group_id="ag-1", - access_group_name="group-1", - access_model_names=["gpt-4", "claude-3"], - access_mcp_server_ids=["mcp-1"], - access_agent_ids=["agent-1"], - ) - ag2 = LiteLLM_AccessGroupTable( - access_group_id="ag-2", - access_group_name="group-2", - access_model_names=["gpt-4", "gemini"], - access_mcp_server_ids=["mcp-1", "mcp-2"], - access_agent_ids=["agent-2"], - ) + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = ["mcp-1"] + row1.access_agent_ids = ["agent-1"] - async def fake_get_access_object(access_group_id, **kwargs): - return {"ag-1": ag1, "ag-2": ag2}[access_group_id] + row2 = MagicMock() + row2.access_group_id = "ag-2" + row2.access_model_names = ["gemini"] + row2.access_mcp_server_ids = ["mcp-2"] + row2.access_agent_ids = ["agent-2"] - with patch( - "litellm.proxy.management_endpoints.team_endpoints._get_access_object", - side_effect=fake_get_access_object, - ): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - MagicMock(), - ): - with patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1", "ag-2"], - ) + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1, row2]) - assert sorted(result["access_group_models"]) == ["claude-3", "gemini", "gpt-4"] - assert sorted(result["access_group_mcp_server_ids"]) == ["mcp-1", "mcp-2"] - assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"]) + + assert result["ag-1"]["models"] == ["gpt-4"] + assert result["ag-2"]["models"] == ["gemini"] @pytest.mark.asyncio - async def test_missing_access_group_skipped(self): - """If an access group doesn't exist, it should be skipped gracefully.""" - from litellm.proxy._types import LiteLLM_AccessGroupTable + async def test_missing_access_group_omitted(self): + """If an access group doesn't exist in DB, it's simply not in the result.""" from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - ag1 = LiteLLM_AccessGroupTable( - access_group_id="ag-1", - access_group_name="group-1", - access_model_names=["gpt-4"], - access_mcp_server_ids=[], - access_agent_ids=[], - ) + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = [] + row1.access_agent_ids = [] - async def fake_get_access_object(access_group_id, **kwargs): - if access_group_id == "ag-1": - return ag1 - raise HTTPException(status_code=404, detail="Not found") + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1]) - with patch( - "litellm.proxy.management_endpoints.team_endpoints._get_access_object", - side_effect=fake_get_access_object, - ): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - MagicMock(), - ): - with patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1", "ag-missing"], - ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-missing"]) - assert result["access_group_models"] == ["gpt-4"] - assert result["access_group_mcp_server_ids"] == [] - assert result["access_group_agent_ids"] == [] + assert "ag-1" in result + assert "ag-missing" not in result @pytest.mark.asyncio - async def test_returns_empty_when_cache_unavailable(self): - """If user_api_key_cache is None, should return empty results.""" + async def test_returns_empty_when_prisma_unavailable(self): + """If prisma_client is None, should return empty dict.""" from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - None, - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1"], - ) + with patch("litellm.proxy.proxy_server.prisma_client", None): + result = await _batch_resolve_access_group_resources(["ag-1"]) - assert result == { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } + assert result == {} + + @pytest.mark.asyncio + async def test_deduplicates_input_ids(self): + """Duplicate IDs in input should result in a single DB lookup.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = [] + row1.access_agent_ids = [] + + fake_find_many = AsyncMock(return_value=[row1]) + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = fake_find_many + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-1", "ag-1"]) + + # Should have been called with deduplicated list + call_args = fake_find_many.call_args + assert len(call_args.kwargs["where"]["access_group_id"]["in"]) == 1 + assert "ag-1" in result From ce219fcc9623bfb9c1584a6014518f764dc2b53d Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 09:36:32 -0700 Subject: [PATCH 21/55] refactor(proxy): extract helpers to fix PLR0915 violations Extract `_apply_non_admin_alias_scope` from `key_aliases`, `_resolve_team_access_group_resources` from `team_info`, and `_enforce_list_team_v2_access` from `list_team_v2` to bring each function under ruff's 50-statement limit. No behavior changes. --- .../key_management_endpoints.py | 65 +++++--- .../management_endpoints/team_endpoints.py | 154 +++++++++++------- 2 files changed, 134 insertions(+), 85 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fe180f19455..497ddcbbe84 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4382,6 +4382,42 @@ async def list_keys( ) +async def _apply_non_admin_alias_scope( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: Any, + query_params: List[Any], + where_parts: List[str], +) -> None: + """Append SQL scope conditions so non-admin users only see aliases for + keys they own or keys belonging to teams they are members of.""" + scope_conditions: List[str] = [] + if user_api_key_dict.user_id: + query_params.append(user_api_key_dict.user_id) + scope_conditions.append(f"user_id = ${len(query_params)}") + + # Look up the user's teams from the user table + user_teams: List[str] = [] + if user_api_key_dict.user_id: + user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id} + ) + if user_row is not None: + user_teams = getattr(user_row, "teams", []) or [] + + if user_teams: + team_placeholders = ", ".join( + f"${len(query_params) + i + 1}" for i in range(len(user_teams)) + ) + query_params.extend(user_teams) + scope_conditions.append(f"team_id IN ({team_placeholders})") + + if scope_conditions: + where_parts.append(f"({' OR '.join(scope_conditions)})") + else: + # No user_id and no teams — return nothing + where_parts.append("FALSE") + + @router.get( "/key/aliases", tags=["key management"], @@ -4442,32 +4478,9 @@ async def key_aliases( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ] if not is_proxy_admin: - scope_conditions: List[str] = [] - if user_api_key_dict.user_id: - query_params.append(user_api_key_dict.user_id) - scope_conditions.append(f"user_id = ${len(query_params)}") - - # Look up the user's teams from the user table - user_teams: List[str] = [] - if user_api_key_dict.user_id: - user_row = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id} - ) - if user_row is not None: - user_teams = getattr(user_row, "teams", []) or [] - - if user_teams: - team_placeholders = ", ".join( - f"${len(query_params) + i + 1}" for i in range(len(user_teams)) - ) - query_params.extend(user_teams) - scope_conditions.append(f"team_id IN ({team_placeholders})") - - if scope_conditions: - where_parts.append(f"({' OR '.join(scope_conditions)})") - else: - # No user_id and no teams — return nothing - where_parts.append("FALSE") + await _apply_non_admin_alias_scope( + user_api_key_dict, prisma_client, query_params, where_parts + ) if search: query_params.append(f"%{search}%") diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae89f09bb44..e1534573789 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2932,6 +2932,25 @@ async def _add_team_member_budget_table( return team_info_response_object +async def _resolve_team_access_group_resources(_team_info: Any) -> None: + """Populate access_group_models / mcp_server_ids / agent_ids on the team + info response by resolving inherited resources from its access groups.""" + if not _team_info.access_group_ids: + return + ag_lookup = await _batch_resolve_access_group_resources( + _team_info.access_group_ids + ) + models, mcp_ids, agent_ids = set(), set(), set() + for ag_id in _team_info.access_group_ids: + if ag_id in ag_lookup: + models.update(ag_lookup[ag_id]["models"]) + mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) + agent_ids.update(ag_lookup[ag_id]["agent_ids"]) + _team_info.access_group_models = list(models) + _team_info.access_group_mcp_server_ids = list(mcp_ids) + _team_info.access_group_agent_ids = list(agent_ids) + + @router.get( "/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @@ -3043,17 +3062,7 @@ async def team_info( ) # Resolve resources inherited from access groups - if _team_info.access_group_ids: - ag_lookup = await _batch_resolve_access_group_resources(_team_info.access_group_ids) - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in _team_info.access_group_ids: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - _team_info.access_group_models = list(models) - _team_info.access_group_mcp_server_ids = list(mcp_ids) - _team_info.access_group_agent_ids = list(agent_ids) + await _resolve_team_access_group_resources(_team_info) response_object = TeamInfoResponseObject( team_id=team_id, @@ -3401,6 +3410,73 @@ def _convert_teams_to_response_models( return team_list +async def _enforce_list_team_v2_access( + user_api_key_dict: UserAPIKeyAuth, + user_id: Optional[str], + organization_id: Optional[str], + prisma_client: Any, + user_api_key_cache: Any, + proxy_logging_obj: Any, +) -> Tuple[Optional[str], Optional[List[str]]]: + """Enforce access control for list_team_v2. + + - Proxy admins and admin viewers can query any teams. + - Org admins can query teams within their organizations. + - Regular users can only query their own teams. + + Returns the (possibly overridden) user_id and org_admin_org_ids. + """ + is_proxy_admin = _user_has_admin_view(user_api_key_dict) + org_admin_org_ids: Optional[List[str]] = None + + if is_proxy_admin: + return user_id, org_admin_org_ids + + # Always check org admin status so that even own-queries see + # the full set of organisation teams, not just direct memberships. + if user_api_key_dict.user_id: + org_admin_org_ids = await _get_org_admin_org_ids( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + if org_admin_org_ids is not None: + # Org admin: validate org_id filter if provided + if organization_id and organization_id not in org_admin_org_ids: + raise HTTPException( + status_code=403, + detail={ + "error": "You can only view teams within your organizations." + }, + ) + verbose_proxy_logger.debug( + "list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s", + user_api_key_dict.user_id, + org_admin_org_ids, + user_id, + ) + else: + # Not an org admin — fall back to standard route check + if not allowed_route_check_inside_route( + user_api_key_dict=user_api_key_dict, requested_user_id=user_id + ): + raise HTTPException( + status_code=401, + detail={ + "error": "Only admin users can query all teams/other teams. Your user role={}".format( + user_api_key_dict.user_role + ) + }, + ) + # Regular user — auto-inject caller's user_id + if user_id is None: + user_id = user_api_key_dict.user_id + + return user_id, org_admin_org_ids + + @router.get( "/v2/team/list", tags=["team management"], @@ -3478,54 +3554,14 @@ async def list_team_v2( ) # --- Access control --- - # Proxy admins and admin viewers can query any teams. - # Org admins can query teams within their organizations. - # Regular users can only query their own teams. - is_proxy_admin = _user_has_admin_view(user_api_key_dict) - org_admin_org_ids: Optional[List[str]] = None - - if not is_proxy_admin: - # Always check org admin status so that even own-queries see - # the full set of organisation teams, not just direct memberships. - if user_api_key_dict.user_id: - org_admin_org_ids = await _get_org_admin_org_ids( - user_id=user_api_key_dict.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - if org_admin_org_ids is not None: - # Org admin: validate org_id filter if provided - if organization_id and organization_id not in org_admin_org_ids: - raise HTTPException( - status_code=403, - detail={ - "error": "You can only view teams within your organizations." - }, - ) - verbose_proxy_logger.debug( - "list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s", - user_api_key_dict.user_id, - org_admin_org_ids, - user_id, - ) - else: - # Not an org admin — fall back to standard route check - if not allowed_route_check_inside_route( - user_api_key_dict=user_api_key_dict, requested_user_id=user_id - ): - raise HTTPException( - status_code=401, - detail={ - "error": "Only admin users can query all teams/other teams. Your user role={}".format( - user_api_key_dict.user_role - ) - }, - ) - # Regular user — auto-inject caller's user_id - if user_id is None: - user_id = user_api_key_dict.user_id + user_id, org_admin_org_ids = await _enforce_list_team_v2_access( + user_api_key_dict=user_api_key_dict, + user_id=user_id, + organization_id=organization_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) if status is not None and status != "deleted": raise HTTPException( From 866c4a25ffb5ec4a7d65e45bc950245a903e5ba9 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 09:47:34 -0700 Subject: [PATCH 22/55] test(ui): update tests to match new team_id / access-group signatures - useKeyAliases, PaginatedKeyAliasSelect: add trailing `undefined` to spy matchers for the new `team_id` param on `useInfiniteKeyAliases` and `keyAliasesCall`. - EntityUsage: mock new `TeamMultiSelect` child so QueryClientProvider is not required for team-entity tests. - ModelsCell: replace the overflow-accordion test with one that verifies the new collapse-on-`all-proxy-models` behavior (no accordion, single badge). --- .../app/(dashboard)/hooks/keys/useKeyAliases.test.ts | 10 +++++----- .../teams/components/TeamsTable/ModelsCell.test.tsx | 12 ++++++------ .../PaginatedKeyAliasSelect.test.tsx | 4 ++-- .../components/EntityUsage/EntityUsage.test.tsx | 4 ++++ 4 files changed, 17 insertions(+), 13 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts index b382b1f2ad3..1e1190b12c8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts @@ -65,7 +65,7 @@ describe("useInfiniteKeyAliases", () => { expect(result.current.isSuccess).toBe(true); }); - expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, undefined); + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, undefined, undefined); expect(result.current.data?.pages[0]).toEqual(mockPage1); }); @@ -74,7 +74,7 @@ describe("useInfiniteKeyAliases", () => { renderHook(() => useInfiniteKeyAliases(25), { wrapper }); await waitFor(() => { - expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 25, undefined); + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 25, undefined, undefined); }); }); @@ -83,7 +83,7 @@ describe("useInfiniteKeyAliases", () => { renderHook(() => useInfiniteKeyAliases(50, "my-alias"), { wrapper }); await waitFor(() => { - expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, "my-alias"); + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, "my-alias", undefined); }); }); @@ -145,7 +145,7 @@ describe("useInfiniteKeyAliases", () => { expect(result.current.data?.pages).toHaveLength(2); }); - expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 2, 2, undefined); + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 2, 2, undefined, undefined); expect(result.current.data?.pages[1]).toEqual(mockPage2); }); @@ -171,7 +171,7 @@ describe("useInfiniteKeyAliases", () => { rerender({ search: "search-result" }); await waitFor(() => { - expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, "search-result"); + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, "search-result", undefined); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.test.tsx index 747ce518cf9..2b487d65322 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.test.tsx @@ -125,14 +125,14 @@ describe("ModelsCell", () => { expect(screen.getByText("+2 more models")).toBeInTheDocument(); }); - it("should render 'all-proxy-models' entries in the overflow section as 'All Proxy Models' badges", () => { + it("should collapse to a single 'All Proxy Models' badge when the models list includes 'all-proxy-models'", () => { renderModelsCell(makeTeam(["m1", "m2", "m3", "all-proxy-models"])); - act(() => { - screen.getByRole("button", { name: /accordion/i }).click(); - }); - - // There should now be an "All Proxy Models" badge in the expanded section + // When all-proxy-models is present, all individual models are hidden and no accordion is shown expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + expect(screen.queryByText("m1")).not.toBeInTheDocument(); + expect(screen.queryByText("m2")).not.toBeInTheDocument(); + expect(screen.queryByText("m3")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /accordion/i })).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.test.tsx b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.test.tsx index 9a3755124b7..79a002cc5a8 100644 --- a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.test.tsx @@ -112,7 +112,7 @@ describe("PaginatedKeyAliasSelect", () => { it("should pass pageSize to useInfiniteKeyAliases", () => { renderWithProviders(); - expect(mockUseInfiniteKeyAliases).toHaveBeenCalledWith(25, undefined); + expect(mockUseInfiniteKeyAliases).toHaveBeenCalledWith(25, undefined, undefined); }); it("should pass search to useInfiniteKeyAliases when user types", async () => { @@ -124,7 +124,7 @@ describe("PaginatedKeyAliasSelect", () => { await user.keyboard("my-alias"); await waitFor(() => { - expect(mockUseInfiniteKeyAliases).toHaveBeenCalledWith(50, "my-alias"); + expect(mockUseInfiniteKeyAliases).toHaveBeenCalledWith(50, "my-alias", undefined); }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx index 5c23cf71ab4..dc201cccfea 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx @@ -45,6 +45,10 @@ vi.mock("../../../EntityUsageExport", () => ({ UsageExportHeader: () =>
Usage Export Header
, })); +vi.mock("../../../common_components/team_multi_select", () => ({ + default: () =>
Team Multi Select
, +})); + // Mock useTeams hook vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: vi.fn(() => ({ From c495acda1b5801690d1b11c96bc22df047a19180 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 09:47:38 -0700 Subject: [PATCH 23/55] fix(ui): send null (not '') for cleared organization_id on team update AntD ({ + value: m, + label: m, + }))} + /> + + + + + + + + remove(name)} + style={{ color: "#ef4444" }} + /> + + ))} + + + + + )} + + + @@ -1162,6 +1281,22 @@ const TeamInfoView: React.FC = ({ Rate Limits
TPM: {info.tpm_limit || "Unlimited"}
RPM: {info.rpm_limit || "Unlimited"}
+ {(() => { + const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record; + const modelRpm = (info.metadata?.model_rpm_limit ?? {}) as Record; + const models = Array.from(new Set([...Object.keys(modelTpm), ...Object.keys(modelRpm)])); + if (models.length === 0) return null; + return ( +
+ Per-model limits: + {models.map((m) => ( +
+ {m}: TPM {modelTpm[m] ?? "—"}, RPM {modelRpm[m] ?? "—"} +
+ ))} +
+ ); + })()}
Team Budget From f0bbd415c8ca5d9ec78e3de73612058ad60b9603 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 13:50:19 -0700 Subject: [PATCH 33/55] fix(ui): require TPM or RPM when adding a per-model team rate limit Previously, a row with a model selected but both limits blank was silently dropped on save (neither model_tpm_limit nor model_rpm_limit got the key), so the row disappeared on reload with no feedback. Now the TPM field's validator blocks submission with "Set at least one of TPM or RPM" when a row has a model but neither limit filled. --- .../src/components/team/TeamInfo.tsx | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 346686d0193..bc54cac9cc7 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1027,7 +1027,21 @@ const TeamInfoView: React.FC = ({ }))} /> - + { + const row = (form.getFieldValue("modelLimits") ?? [])[name] ?? {}; + if (row.model && value == null && row.rpm == null) { + return Promise.reject(new Error("Set at least one of TPM or RPM")); + } + return Promise.resolve(); + }, + }, + ]} + > From 9dca4319892642650644013f6fcd4185ed8947eb Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 31 Mar 2026 10:02:42 -0700 Subject: [PATCH 34/55] fix(ui): use entity key for export instead of extracting team_id from api_key_breakdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The export utility always extracted team_id from api_key_breakdown metadata to populate the entity label/ID columns. This worked for team exports (where entity key = team_id) but was wrong for every other entity type — tags, orgs, customers, agents, users all showed the API key's team name (e.g. "admins") instead of the actual entity value. Replace extractTeamIdFromApiKeyBreakdown with resolveEntityDisplay which uses the entity key directly. For teams the teamAliasMap still resolves a human-readable alias; for all other types the entity key itself is the correct label. --- .../src/components/EntityUsageExport/utils.ts | 80 +++++++------------ 1 file changed, 30 insertions(+), 50 deletions(-) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 45bf21a6e7d..6013e276481 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -3,19 +3,16 @@ import type { DateRangePickerValue } from "@tremor/react"; import Papa from "papaparse"; import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; -// Helper function to extract team_id from api_key_breakdown -const extractTeamIdFromApiKeyBreakdown = (apiKeyBreakdown: Record | undefined): string | null => { - if (!apiKeyBreakdown) return null; - - // Look through all API keys to find the first non-null team_id - for (const apiKeyData of Object.values(apiKeyBreakdown)) { - const teamId = (apiKeyData as any)?.metadata?.team_id; - if (teamId) { - return teamId; - } - } - return null; -}; +// Resolve display name for an entity. For teams the teamAliasMap provides +// a human-readable alias; for every other entity type the entity key itself +// (tag name, org id, customer id, …) is already the correct label. +const resolveEntityDisplay = ( + entity: string, + teamAliasMap: Record, +): { id: string; alias: string } => ({ + id: entity, + alias: teamAliasMap[entity] || entity, +}); // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). // If the backend adds a field, add it here too. @@ -68,18 +65,7 @@ export const getEntityBreakdown = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty) - const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown) || entity; - // Extract key_alias from the first API key that has one - const apiKeyBreakdown = data.api_key_breakdown || {}; - let keyAlias: string | null = null; - for (const apiKeyData of Object.values(apiKeyBreakdown)) { - const alias = (apiKeyData as any)?.metadata?.key_alias; - if (alias) { - keyAlias = alias; - break; - } - } + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); if (!entitySpend[entity]) { entitySpend[entity] = { @@ -95,8 +81,8 @@ export const getEntityBreakdown = ( cache_creation_input_tokens: 0, }, metadata: { - alias: keyAlias || teamAliasMap[teamId] || entity, - id: teamId, + alias, + id, }, }; } @@ -124,14 +110,12 @@ export const generateDailyData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty) - const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown); - const teamAlias = teamId ? teamAliasMap[teamId] || null : null; + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); dailyBreakdown.push({ Date: day.date, - [entityLabel]: teamAlias || "-", - [`${entityLabel} ID`]: teamId || "-", + [entityLabel]: alias, + [`${entityLabel} ID`]: id, "Spend ($)": formatNumberWithCommas(data.metrics.spend, 4), Requests: data.metrics.api_requests, "Successful Requests": data.metrics.successful_requests, @@ -151,12 +135,12 @@ export const generateDailyWithKeysData = ( entityLabel: string, teamAliasMap: Record = {}, ): any[] => { - // Aggregate by unique (Date, Team ID, Key ID) combination to prevent duplicates + // Aggregate by unique (Date, Entity ID, Key ID) combination to prevent duplicates const aggregatedData: { [key: string]: { Date: string; - teamId: string; - teamAlias: string | null; + entityId: string; + entityAlias: string; keyId: string; keyAlias: string | null; metrics: { @@ -173,23 +157,22 @@ export const generateDailyWithKeysData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { + const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap); const apiKeyBreakdown = data.api_key_breakdown || {}; // Iterate through each API key in the breakdown Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => { const keyAlias = keyData?.metadata?.key_alias || null; - const teamId = keyData?.metadata?.team_id || entity; - const teamAlias = teamId ? teamAliasMap[teamId] || null : null; - // Create unique key for aggregation: Date_TeamID_KeyID - const uniqueKey = `${day.date}_${teamId}_${keyId}`; + // Create unique key for aggregation: Date_EntityID_KeyID + const uniqueKey = `${day.date}_${entityId}_${keyId}`; if (!aggregatedData[uniqueKey]) { - // First time seeing this (Date, Team ID, Key ID) combination + // First time seeing this (Date, Entity ID, Key ID) combination aggregatedData[uniqueKey] = { Date: day.date, - teamId, - teamAlias, + entityId, + entityAlias, keyId, keyAlias, metrics: { @@ -219,8 +202,8 @@ export const generateDailyWithKeysData = ( // Convert aggregated data to array format const dailyKeyBreakdown = Object.values(aggregatedData).map((item) => ({ Date: item.Date, - [entityLabel]: item.teamAlias || "-", - [`${entityLabel} ID`]: item.teamId || "-", + [entityLabel]: item.entityAlias, + [`${entityLabel} ID`]: item.entityId, "Key Alias": item.keyAlias || "-", "Key ID": item.keyId, "Spend ($)": formatNumberWithCommas(item.metrics.spend, 4), @@ -273,16 +256,13 @@ export const generateDailyWithModelsData = ( }); Object.entries(dailyEntityModels).forEach(([entity, models]) => { - const entityData = resolveEntities(day.breakdown)[entity]; - // Extract team_id from api_key_breakdown metadata (not entityData.metadata which is empty) - const teamId = extractTeamIdFromApiKeyBreakdown(entityData?.api_key_breakdown); - const teamAlias = teamId ? teamAliasMap[teamId] || null : null; + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); Object.entries(models).forEach(([model, metrics]: [string, any]) => { dailyModelBreakdown.push({ Date: day.date, - [entityLabel]: teamAlias || "-", - [`${entityLabel} ID`]: teamId || "-", + [entityLabel]: alias, + [`${entityLabel} ID`]: id, Model: model, "Spend ($)": formatNumberWithCommas(metrics.spend, 4), Requests: metrics.requests, From 08df8643bf00d45d5e82c8f07c4c167056e70237 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Sun, 5 Apr 2026 02:34:31 +0530 Subject: [PATCH 35/55] fix(docker): include enterprise bridge in non-root runtime image (#24917) Copy the /app/enterprise bridge package into the non-root runtime image so enterprise proxy hooks register correctly (including managed_files). --- docker/Dockerfile.non_root | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index bf1af9ed756..8e911e95ffa 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -142,6 +142,9 @@ COPY --from=builder /app/requirements.txt /app/requirements.txt COPY --from=builder /app/docker/entrypoint.sh /app/docker/prod_entrypoint.sh /app/docker/ COPY --from=builder /app/docker/supervisord.conf /etc/supervisord.conf COPY --from=builder /app/schema.prisma /app/ +# Keep enterprise bridge module in runtime so `enterprise.enterprise_hooks` +# can load and register managed enterprise hooks (e.g. managed_files). +COPY --from=builder /app/enterprise /app/enterprise # Copy prisma_migration.py for Helm migrations job compatibility COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py COPY --from=builder /wheels/ /wheels/ From 8e6300f0bf501b57e431d41099e1cb0686e5f8f1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 4 Apr 2026 14:25:19 -0700 Subject: [PATCH 36/55] fix: replace unsupported Prisma JSON path filter in _get_team_deployments prisma-client-py 0.11.0 does not support JSON path filtering (path/equals) on Json fields. Replace with model_name prefix query + Python-side team_id confirmation. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../model_management_endpoints.py | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 2ca8e3daba3..754727d4716 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -485,18 +485,31 @@ async def _get_team_deployments( Centralizes team deployment queries to ensure consistent filtering and error handling. This is the established helper pattern for team deployment DB access in this module. - Note: Direct Prisma call is intentional here as this IS the helper function that - encapsulates the DB access pattern for team deployments. + Note: prisma-client-py 0.11.0 does not support JSON path filtering, so we filter + by the model_name prefix (team models use "model_name_{team_id}_*") and confirm + team_id in model_info with Python-side filtering. """ + prefix = f"model_name_{team_id}_" response = await prisma_client.db.litellm_proxymodeltable.find_many( where={ - "model_info": { - "path": ["team_id"], - "equals": team_id, - } + "model_name": {"startswith": prefix}, } ) - return response if response else [] + if not response: + return [] + + # Confirm team_id in model_info (defensive check) + result = [] + for row in response: + model_info = row.model_info + if isinstance(model_info, str): + try: + model_info = json.loads(model_info) + except (TypeError, ValueError): + continue + if isinstance(model_info, dict) and model_info.get("team_id") == team_id: + result.append(row) + return result async def _update_existing_team_model_assignment( From 566a04126fb1091790fd487956b7f6a26d445c60 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 4 Apr 2026 14:34:39 -0700 Subject: [PATCH 37/55] test: add unit tests for _get_team_deployments filtering logic Tests cover: matching deployments, wrong team_id filtering, string-encoded model_info, empty results, invalid model_info, and mixed deployment filtering. Also updates MockPrismaClient.find_many to support the new startswith query. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../test_model_management_endpoints.py | 131 +++++++++++++----- 1 file changed, 98 insertions(+), 33 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index c3e3d1ecbd9..198cd39fca0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, + _get_team_deployments, clear_cache, ) from litellm.proxy.utils import PrismaClient @@ -53,44 +54,23 @@ class MockPrismaClient: ) return None - async def find_many(self, where): - # Filter sibling deployments by team_id if where clause specifies it + async def find_many(self, where=None): + # Filter sibling deployments based on where clause if not self.sibling_deployments: return [] - # Extract team_id from where clause if present - team_id_filter = None - if where and "model_info" in where: - model_info_filter = where["model_info"] - if isinstance(model_info_filter, dict) and "path" in model_info_filter: - if ( - model_info_filter["path"] == ["team_id"] - and "equals" in model_info_filter - ): - team_id_filter = model_info_filter["equals"] + results = self.sibling_deployments - # Filter deployments by team_id if specified - if team_id_filter: + # Support model_name startswith filter (used by _get_team_deployments) + if where and "model_name" in where: + model_name_filter = where["model_name"] + if isinstance(model_name_filter, dict) and "startswith" in model_name_filter: + prefix = model_name_filter["startswith"] + results = [ + d for d in results if d.model_name.startswith(prefix) + ] - def _get_team_id(model_info): - if isinstance(model_info, dict): - return model_info.get("team_id") - if isinstance(model_info, str): - try: - parsed = json.loads(model_info) - except (TypeError, ValueError): - return None - if isinstance(parsed, dict): - return parsed.get("team_id") - return None - - return [ - d - for d in self.sibling_deployments - if _get_team_id(d.model_info) == team_id_filter - ] - - return self.sibling_deployments + return results @property def litellm_teamtable(self): @@ -1276,3 +1256,88 @@ class TestAddAndDeleteModelLifecycle: user_api_key_dict=admin_user, ) assert str(exc_info.value.code) == "400" + + +class TestGetTeamDeployments: + """Tests for _get_team_deployments which filters by model_name prefix + Python-side team_id check.""" + + @pytest.mark.asyncio + async def test_returns_matching_team_deployments(self): + """Deployments with matching model_name prefix and team_id are returned.""" + team_id = "team_abc" + dep = MagicMock() + dep.model_name = f"model_name_{team_id}_uuid1" + dep.model_info = {"team_id": team_id, "team_public_model_name": "gpt-4"} + + prisma_client = MockPrismaClient(sibling_deployments=[dep]) + result = await _get_team_deployments(team_id, prisma_client) + assert len(result) == 1 + assert result[0] is dep + + @pytest.mark.asyncio + async def test_filters_out_wrong_team_id_in_model_info(self): + """A deployment whose model_name matches but model_info.team_id differs is excluded.""" + team_id = "team_abc" + dep = MagicMock() + dep.model_name = f"model_name_{team_id}_uuid1" + dep.model_info = {"team_id": "other_team"} + + prisma_client = MockPrismaClient(sibling_deployments=[dep]) + result = await _get_team_deployments(team_id, prisma_client) + assert len(result) == 0 + + @pytest.mark.asyncio + async def test_handles_string_encoded_model_info(self): + """Legacy rows with JSON-string model_info are parsed and filtered correctly.""" + team_id = "team_abc" + dep = MagicMock() + dep.model_name = f"model_name_{team_id}_uuid1" + dep.model_info = json.dumps({"team_id": team_id}) + + prisma_client = MockPrismaClient(sibling_deployments=[dep]) + result = await _get_team_deployments(team_id, prisma_client) + assert len(result) == 1 + + @pytest.mark.asyncio + async def test_returns_empty_when_no_deployments(self): + """Returns empty list when no deployments exist.""" + prisma_client = MockPrismaClient(sibling_deployments=[]) + result = await _get_team_deployments("team_abc", prisma_client) + assert result == [] + + @pytest.mark.asyncio + async def test_skips_rows_with_invalid_model_info(self): + """Rows with non-dict, non-parseable model_info are skipped.""" + team_id = "team_abc" + dep = MagicMock() + dep.model_name = f"model_name_{team_id}_uuid1" + dep.model_info = "not-valid-json" + + prisma_client = MockPrismaClient(sibling_deployments=[dep]) + result = await _get_team_deployments(team_id, prisma_client) + assert len(result) == 0 + + @pytest.mark.asyncio + async def test_multiple_deployments_mixed_filtering(self): + """Only deployments with correct prefix AND team_id are returned.""" + team_id = "team_abc" + + # Matches both prefix and team_id + dep1 = MagicMock() + dep1.model_name = f"model_name_{team_id}_uuid1" + dep1.model_info = {"team_id": team_id} + + # Matches prefix but wrong team_id + dep2 = MagicMock() + dep2.model_name = f"model_name_{team_id}_uuid2" + dep2.model_info = {"team_id": "wrong_team"} + + # Different prefix entirely (won't be returned by mock's startswith filter) + dep3 = MagicMock() + dep3.model_name = "model_name_other_team_uuid3" + dep3.model_info = {"team_id": "other_team"} + + prisma_client = MockPrismaClient(sibling_deployments=[dep1, dep2, dep3]) + result = await _get_team_deployments(team_id, prisma_client) + assert len(result) == 1 + assert result[0] is dep1 From 5452692af4bce8d0af456125749b27ea7f9c86ee Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 14:37:38 -0700 Subject: [PATCH 38/55] test(ui): align EntityUsageExport tests with entity-key display logic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fix (9dca431989) switched the export display logic to use the entity key directly, but left the unit tests asserting the old behavior where id/alias were extracted from api_key_breakdown metadata. Rename mock entity keys from entity1/entity2 to "team-1"/"team-2" to reflect the real shape of team-export payloads (breakdown.entities is keyed by the entity identifier). Replace three tests whose assertions documented the removed behavior: - "should use key alias when available" → entity key is the alias when no team alias map is supplied - "should use team alias map when key alias is not available" → team alias map resolves from the entity key - "should use dash when team id is not available" → entity key itself is the fallback label (illustrated with a tag export) --- .../EntityUsageExport/utils.test.ts | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 2ca7ad7ef31..a462a9f53d1 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -30,13 +30,17 @@ vi.mock("papaparse", () => ({ })); describe("EntityUsageExport utils", () => { + // Entity keys match team_ids because that's how the backend shapes team exports + // (breakdown.entities is keyed by team_id). The fix under test uses the entity key + // directly for display, so the key_alias/team_id in api_key_breakdown metadata is + // no longer consulted — it's retained here only to mirror real payload shape. const mockSpendData: EntitySpendData = { results: [ { date: "2025-01-01", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 10.5, api_requests: 100, @@ -64,7 +68,7 @@ describe("EntityUsageExport utils", () => { }, }, }, - entity2: { + "team-2": { metrics: { spend: 20.3, api_requests: 200, @@ -99,7 +103,7 @@ describe("EntityUsageExport utils", () => { date: "2025-01-02", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 15.2, api_requests: 150, @@ -184,14 +188,16 @@ describe("EntityUsageExport utils", () => { expect(entity1?.metrics.cache_creation_input_tokens).toBe(75); }); - it("should use key alias when available", () => { + it("should use entity key as alias when no team alias map is provided", () => { + // Non-team exports (tags, orgs, customers, …) pass no teamAliasMap. + // For teams, this is also the fallback when a team is missing from the map. const result = getEntityBreakdown(mockSpendData); const entity1 = result.find((e) => e.metadata.id === "team-1"); - expect(entity1?.metadata.alias).toBe("alias-1"); + expect(entity1?.metadata.alias).toBe("team-1"); }); - it("should use team alias map when key alias is not available", () => { + it("should use team alias map to resolve alias from entity key", () => { const spendDataWithoutAlias: EntitySpendData = { ...mockSpendData, results: [ @@ -199,7 +205,7 @@ describe("EntityUsageExport utils", () => { date: "2025-01-01", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 10.5, api_requests: 100, @@ -299,7 +305,7 @@ describe("EntityUsageExport utils", () => { date: "2025-01-01", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 10.5, api_requests: 100, @@ -379,15 +385,17 @@ describe("EntityUsageExport utils", () => { } }); - it("should use dash when team id is not available", () => { - const spendDataWithoutTeamId: EntitySpendData = { + it("should fall back to the entity key when there is no team alias mapping", () => { + // e.g. tag/org/customer exports where teamAliasMap has no entry for the entity, + // or a team that isn't in the alias map — the entity key itself is the label. + const spendDataWithoutAlias: EntitySpendData = { ...mockSpendData, results: [ { date: "2025-01-01", breakdown: { entities: { - entity1: { + "my-tag": { metrics: { spend: 10.5, api_requests: 100, @@ -406,11 +414,11 @@ describe("EntityUsageExport utils", () => { metadata: mockSpendData.metadata, }; - const result = generateDailyData(spendDataWithoutTeamId, "Team"); + const result = generateDailyData(spendDataWithoutAlias, "Tag"); const entry = result[0]; - expect(entry["Team ID"]).toBe("-"); - expect(entry["Team"]).toBe("-"); + expect(entry["Tag ID"]).toBe("my-tag"); + expect(entry["Tag"]).toBe("my-tag"); }); it("should format spend values correctly", () => { @@ -471,7 +479,7 @@ describe("EntityUsageExport utils", () => { date: "2025-01-01", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 10.5, api_requests: 100, @@ -514,7 +522,7 @@ describe("EntityUsageExport utils", () => { }, }, }, - entity2: { + "team-2": { metrics: { spend: 20.3, api_requests: 200, @@ -549,7 +557,7 @@ describe("EntityUsageExport utils", () => { date: "2025-01-02", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 15.2, api_requests: 150, @@ -979,7 +987,7 @@ describe("EntityUsageExport utils", () => { date: "2025-01-01", breakdown: { entities: { - entity1: { + "team-1": { metrics: { spend: 10.5, api_requests: 100, From b53cfe729abc3210b245a204bf2577713704dd85 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Sat, 4 Apr 2026 14:44:07 -0700 Subject: [PATCH 39/55] Litellm ishaan march30 (#24887) (#25151) * fix(pricing): add unversioned vertex_ai/claude-haiku-4-5 entry Missing unversioned entry causes cost tracking to return $0.00 for all requests using vertex_ai/claude-haiku-4-5. All other Vertex AI Claude models have both versioned and unversioned entries. * fix(router): skip misleading tags error when no candidates (e.g. cooldown) Return early from get_deployments_for_tag when healthy_deployments is empty so tag-based routing does not raise no_deployments_with_tag_routing after cooldown filters all deployments. Adds regression test. Made-with: Cursor * feat(oci): add embedding support and update model catalog - Add OCIEmbeddingConfig for OCI GenAI embedding models - Add 16 new chat models (Cohere, Meta Llama, xAI Grok, Google Gemini) - Add 8 embedding models (Cohere embed v3.0, v4.0) - Update documentation with embedding examples - Update pricing for all new models * test(oci): add unit tests for OCI embedding support - 17 unit tests covering OCIEmbeddingConfig - Tests for URL generation, param mapping, request/response transform - Tests for model pricing JSON completeness * style(oci): format with black and ruff * fix(oci): correct embedding request body format OCI embedText API expects inputs, truncate, and inputType at the top level of the request body, not nested under embedTextDetails. Fixed transformation and updated tests accordingly. Verified with real OCI API: 3/3 embedding models working. * docs: clarify tag routing early return and test intent Made-with: Cursor * fix(oci): address code review findings from Greptile - P1: Fix signing URL mismatch with custom api_base by accepting api_base parameter in transform_embedding_request - P2: Remove encoding_format from supported params (OCI does not support it, was silently dropped) - P2: Raise ValueError for token-array inputs instead of silently converting to string representation - Add test for token-list rejection * fix(mcp): add STS AssumeRole support for MCP SigV4 authentication MCPSigV4Auth only supported static AWS credentials or the boto3 default credential chain. Production Kubernetes environments typically authenticate via IAM role assumption (sts:AssumeRole), which was not possible. Add aws_role_name and aws_session_name parameters to the MCP SigV4 auth stack. When aws_role_name is provided, MCPSigV4Auth calls sts:AssumeRole to obtain temporary credentials before signing requests. Explicit keys, if also provided, are used as the source identity for the STS call; otherwise ambient credentials (pod role, instance profile) are used. * fix: stop logging credential values and add missing redaction patterns Replaces raw credential values in debug/error log messages with boolean presence checks or type names. Adds PEM block, GCP token, JWT, SAS token, and service-account blob patterns to the redaction filter. Fixes private_key pattern to capture full PEM blocks instead of stopping at the first whitespace. Addresses: Vertex AI credential JSON (including RSA private key) being logged to stderr on health check failures. * fix: log only field names for UserAPIKeyAuth, not full object * style: apply black formatting to experimental_mcp_client/client.py * style: fix black/isort formatting and mypy error in proxy_server.py - Fix black formatting in experimental_mcp_client/client.py (done in prev commit) - Fix black/isort formatting in key_management_endpoints.py, proxy_server.py, transformation.py - Fix mypy: iterate over optional list safely (access_group_ids or []) in proxy_server.py * fix(test): patch check_migration.verbose_logger directly to fix xdist ordering issue When test_proxy_cli.py tests run before test_check_migration.py in the same xdist worker, litellm.proxy.db.check_migration is already in sys.modules. Patching litellm._logging.verbose_logger has no effect on the already-bound reference. Patch the correct target (check_migration.verbose_logger) and import the module before patching so the order doesn't matter. * fix(mypy): make api_base Optional in PydanticAIProviderConfig to match base class signature --------- Co-authored-by: Ihsan Soydemir Co-authored-by: Milan Co-authored-by: Daniel Gandolfi Co-authored-by: Claude Sonnet 4.6 Co-authored-by: michelligabriele Co-authored-by: user <70670632+stuxf@users.noreply.github.com> Co-authored-by: Ishaan Jaffer --- docs/my-website/docs/mcp.md | 3 +- docs/my-website/docs/mcp_aws_sigv4.md | 53 ++- docs/my-website/docs/providers/oci.md | 109 +++++- litellm/__init__.py | 1 + litellm/_logging.py | 38 +- .../providers/pydantic_ai_agents/config.py | 10 +- litellm/experimental_mcp_client/client.py | 50 ++- .../azure_storage/azure_storage.py | 5 +- .../gcs_bucket/gcs_bucket_base.py | 8 +- .../get_llm_provider_logic.py | 4 +- litellm/llms/azure/common_utils.py | 14 +- litellm/llms/bedrock/base_aws_llm.py | 16 +- litellm/llms/oci/chat/transformation.py | 9 +- litellm/llms/oci/embed/__init__.py | 0 litellm/llms/oci/embed/transformation.py | 347 ++++++++++++++++ litellm/llms/vertex_ai/vertex_llm_base.py | 28 +- litellm/main.py | 16 + ...odel_prices_and_context_window_backup.json | 305 ++++++++++++++- .../mcp_server/mcp_server_manager.py | 8 + litellm/proxy/auth/oauth2_check.py | 4 +- litellm/proxy/auth/oauth2_proxy_hook.py | 8 +- .../proxy/hooks/parallel_request_limiter.py | 2 +- .../key_management_endpoints.py | 10 +- litellm/proxy/proxy_server.py | 27 +- .../spend_management_endpoints.py | 4 +- litellm/proxy/utils.py | 2 +- .../transformation.py | 6 +- litellm/router_strategy/tag_based_routing.py | 7 + litellm/router_utils/handle_error.py | 5 +- .../secret_managers/secret_manager_handler.py | 6 +- litellm/types/mcp.py | 6 + .../types/mcp_server/mcp_server_manager.py | 2 + litellm/utils.py | 4 + model_prices_and_context_window.json | 305 ++++++++++++++- tests/test_litellm/llms/oci/embed/__init__.py | 0 .../llms/oci/embed/test_oci_embedding.py | 369 ++++++++++++++++++ .../mcp_server/test_mcp_sigv4_auth.py | 214 ++++++++++ .../proxy/db/test_check_migration.py | 10 +- .../test_router_tag_regex_routing.py | 20 + tests/test_litellm/test_secret_redaction.py | 77 ++++ .../mcp_tools/create_mcp_server.tsx | 32 ++ .../components/mcp_tools/mcp_server_edit.tsx | 32 ++ 42 files changed, 2083 insertions(+), 93 deletions(-) create mode 100644 litellm/llms/oci/embed/__init__.py create mode 100644 litellm/llms/oci/embed/transformation.py create mode 100644 tests/test_litellm/llms/oci/embed/__init__.py create mode 100644 tests/test_litellm/llms/oci/embed/test_oci_embedding.py diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index b805cce4d7a..f6fe01ac28f 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -278,7 +278,8 @@ mcp_servers: url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" transport: "http" auth_type: "aws_sigv4" - aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_role_name: os.environ/AWS_ROLE_ARN # optional — IAM role to assume + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # optional — falls back to IAM role aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY aws_region_name: us-east-1 aws_service_name: bedrock-agentcore diff --git a/docs/my-website/docs/mcp_aws_sigv4.md b/docs/my-website/docs/mcp_aws_sigv4.md index 9dc60bce06e..e556ad244f8 100644 --- a/docs/my-website/docs/mcp_aws_sigv4.md +++ b/docs/my-website/docs/mcp_aws_sigv4.md @@ -36,6 +36,8 @@ LiteLLM's `aws_sigv4` auth type handles this automatically: every outgoing MCP r | **AWS Access Key ID** | No | Falls back to boto3 credential chain if blank | | **AWS Secret Access Key** | No | Required if Access Key ID is provided | | **AWS Session Token** | No | Only needed for temporary STS credentials | +| **AWS Role ARN** | No | IAM role ARN for STS AssumeRole (e.g., `arn:aws:iam::123456789012:role/MyRole`). If set, LiteLLM assumes this role before signing | +| **AWS Session Name** | No | Session name for the AssumeRole call — appears in CloudTrail. Auto-generated if omitted | Once created, LiteLLM will sign every outgoing MCP request with SigV4. The server's tools appear automatically in the MCP Tools list. @@ -66,8 +68,8 @@ mcp_servers: url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" transport: "http" auth_type: "aws_sigv4" - aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID - aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_role_name: os.environ/AWS_ROLE_ARN # IAM role to assume (recommended) + aws_session_name: "litellm-prod" # optional — for CloudTrail auditing aws_region_name: "us-east-1" aws_service_name: "bedrock-agentcore" ``` @@ -128,6 +130,8 @@ curl http://localhost:4000/mcp-rest/tools/call \ | `aws_region_name` | Yes | AWS region (e.g., `us-east-1`) | | `aws_service_name` | No | AWS service name for signing. Defaults to `bedrock-agentcore` | | `aws_session_token` | No | AWS session token for temporary credentials. Supports `os.environ/VAR_NAME` | +| `aws_role_name` | No | IAM role ARN for STS AssumeRole. Supports `os.environ/VAR_NAME`. When set, LiteLLM calls `sts:AssumeRole` to get temporary credentials before signing | +| `aws_session_name` | No | Session name for the AssumeRole call (appears in CloudTrail). Auto-generated if omitted. Supports `os.environ/VAR_NAME` | ## How It Works @@ -157,6 +161,42 @@ mcp_servers: aws_service_name: "bedrock-agentcore" ``` +## Using IAM Role Assumption (AssumeRole) + +For production environments where your LiteLLM instance authenticates via an IAM role (e.g., EKS pod role, EC2 instance profile), you can configure `aws_role_name` to have LiteLLM call `sts:AssumeRole` before signing MCP requests: + +```yaml title="config.yaml with AssumeRole" showLineNumbers +mcp_servers: + my_agentcore_mcp: + url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" + transport: "http" + auth_type: "aws_sigv4" + aws_role_name: "arn:aws:iam::123456789012:role/BedrockAgentCoreRole" + aws_session_name: "litellm-prod" # optional + aws_region_name: "us-east-1" + aws_service_name: "bedrock-agentcore" +``` + +LiteLLM uses the ambient credentials (pod role, instance profile, or env vars) to call `sts:AssumeRole`, then signs MCP requests with the assumed role's temporary credentials. + +You can also combine `aws_role_name` with explicit access keys — the keys are then used as the source identity for the AssumeRole call: + +```yaml title="config.yaml with AssumeRole + explicit source keys" showLineNumbers +mcp_servers: + my_agentcore_mcp: + url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" + transport: "http" + auth_type: "aws_sigv4" + aws_role_name: os.environ/AWS_ROLE_ARN + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: "us-east-1" +``` + +:::tip +For most Kubernetes deployments, you only need `aws_role_name` and `aws_region_name` — the pod's IAM role provides the source credentials automatically. +::: + ## Troubleshooting ### 403 Forbidden from AWS @@ -166,6 +206,15 @@ mcp_servers: - Ensure `aws_service_name` is set to `bedrock-agentcore` - If using STS credentials, confirm `aws_session_token` is set and not expired +### AssumeRole AccessDenied + +If you get `AccessDenied` when using `aws_role_name`: + +- Verify the role ARN is correct +- Check that the trust policy on the target role allows your source identity to assume it +- If running on EKS, ensure the pod's service account is annotated with the correct IAM role +- Check CloudTrail for the failed `sts:AssumeRole` call to see the exact error + ### Health check errors on startup SigV4-authenticated MCP servers skip the standard health check on proxy startup. This is expected — the proxy will still sign requests correctly when tools are invoked. diff --git a/docs/my-website/docs/providers/oci.md b/docs/my-website/docs/providers/oci.md index ce6fe18dd6f..1d7a0a3d502 100644 --- a/docs/my-website/docs/providers/oci.md +++ b/docs/my-website/docs/providers/oci.md @@ -8,24 +8,54 @@ Check the [OCI Models List](https://docs.oracle.com/en-us/iaas/Content/generativ ## Supported Models -### Meta Llama Models +### Chat / Text Generation + +#### Meta Llama Models - `meta.llama-4-maverick-17b-128e-instruct-fp8` - `meta.llama-4-scout-17b-16e-instruct` - `meta.llama-3.3-70b-instruct` +- `meta.llama-3.3-70b-instruct-fp8-dynamic` - `meta.llama-3.2-90b-vision-instruct` +- `meta.llama-3.2-11b-vision-instruct` - `meta.llama-3.1-405b-instruct` +- `meta.llama-3.1-70b-instruct` -### xAI Grok Models +#### xAI Grok Models +- `xai.grok-4.20` +- `xai.grok-4.20-multi-agent` - `xai.grok-4` +- `xai.grok-4-fast` +- `xai.grok-4.1-fast` - `xai.grok-3` - `xai.grok-3-fast` - `xai.grok-3-mini` - `xai.grok-3-mini-fast` +- `xai.grok-code-fast-1` -### Cohere Models +#### Cohere Models - `cohere.command-latest` - `cohere.command-a-03-2025` +- `cohere.command-a-reasoning-08-2025` +- `cohere.command-a-vision-07-2025` +- `cohere.command-a-translate-08-2025` - `cohere.command-plus-latest` +- `cohere.command-r-08-2024` +- `cohere.command-r-plus-08-2024` + +#### Google Gemini Models (via OCI) +- `google.gemini-2.5-pro` +- `google.gemini-2.5-flash` +- `google.gemini-2.5-flash-lite` + +### Embedding Models +- `cohere.embed-english-v3.0` (1024 dimensions) +- `cohere.embed-english-light-v3.0` (384 dimensions) +- `cohere.embed-multilingual-v3.0` (1024 dimensions) +- `cohere.embed-multilingual-light-v3.0` (384 dimensions) +- `cohere.embed-english-image-v3.0` (1024 dimensions, multimodal) +- `cohere.embed-english-light-image-v3.0` (384 dimensions, multimodal) +- `cohere.embed-multilingual-light-image-v3.0` (384 dimensions, multimodal) +- `cohere.embed-v4.0` (1536 dimensions, multimodal) ## Authentication @@ -394,4 +424,75 @@ response = completion( | `oci_tenancy` | string | - | (Manual auth) The OCID of your OCI tenancy | | `oci_key` | string | - | (Manual auth) The private key content as a string | | `oci_key_file` | string | - | (Manual auth) Path to the private key file | -| `oci_signer` | object | - | (SDK auth) OCI SDK Signer object for authentication | \ No newline at end of file +| `oci_signer` | object | - | (SDK auth) OCI SDK Signer object for authentication | + +## Embeddings + +LiteLLM supports OCI Generative AI embedding models. These models use the same authentication methods described above. + + + + +```python +from litellm import embedding + +response = embedding( + model="oci/cohere.embed-english-v3.0", + input=["Hello world", "Goodbye world"], + oci_region="us-ashburn-1", + oci_user=, + oci_fingerprint=, + oci_tenancy=, + oci_key=, + oci_compartment_id=, +) +print(response) +``` + + + + +```python +from litellm import embedding +from oci.signer import Signer + +signer = Signer( + tenancy="ocid1.tenancy.oc1..", + user="ocid1.user.oc1..", + fingerprint="xx:xx:xx:xx:xx:xx:xx:xx:xx:xx:xx:xx:xx:xx:xx:xx", + private_key_file_location="~/.oci/key.pem", +) + +response = embedding( + model="oci/cohere.embed-english-v3.0", + input=["Hello world", "Goodbye world"], + oci_signer=signer, + oci_region="us-ashburn-1", + oci_compartment_id="", +) +print(response) +``` + + + + +### Embedding Optional Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input_type` | string | - | The type of input: `search_document`, `search_query`, `classification`, `clustering` | +| `truncate` | string | `END` | Truncation strategy when input exceeds max tokens: `END` or `START` | + +### Using Dedicated Embedding Endpoints + +```python +response = embedding( + model="oci/cohere.embed-english-v3.0", + input=["Hello world"], + oci_serving_mode="DEDICATED", + oci_endpoint_id="ocid1.generativeaiendpoint.oc1...", + oci_region="us-ashburn-1", + oci_compartment_id="", + # ... auth params +) +``` \ No newline at end of file diff --git a/litellm/__init__.py b/litellm/__init__.py index e45d926e8db..d4418c661a3 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1838,6 +1838,7 @@ if TYPE_CHECKING: ) from .llms.v0.chat.transformation import V0ChatConfig as V0ChatConfig from .llms.oci.chat.transformation import OCIChatConfig as OCIChatConfig + from .llms.oci.embed.transformation import OCIEmbeddingConfig as OCIEmbeddingConfig from .llms.morph.chat.transformation import MorphChatConfig as MorphChatConfig from .llms.ragflow.chat.transformation import RAGFlowConfig as RAGFlowConfig from .llms.lambda_ai.chat.transformation import ( diff --git a/litellm/_logging.py b/litellm/_logging.py index 65e6045b0b1..62283f6f65a 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -26,6 +26,12 @@ _REDACTED = "REDACTED" def _build_secret_patterns() -> re.Pattern: patterns: List[str] = [ + # ── PEM private key / certificate blocks ── + r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----", + # ── GCP OAuth2 access tokens (ya29.*) ── + r"\bya29\.[A-Za-z0-9_.~+/-]+", + # ── Credential %s formatting (space separator, no key= prefix) ── + r"(?:client_secret|azure_password|azure_username)\s+[^\s,'\"})\]{}>]+", # AWS access key IDs r"(?:AKIA|ASIA)[0-9A-Z]{16}", # AWS secrets / session tokens / access key IDs (key=value) @@ -46,7 +52,8 @@ def _build_secret_patterns() -> re.Pattern: # Google API keys r"AIza[0-9A-Za-z\-_]{35}", # Password / secret params (handles key=value and 'key': 'value') - r"\w*(?:password|passwd|client_secret|secret_key|_secret)" + # Word boundary prevents O(n^2) backtracking on long word-char runs. + r"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)" r"['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+", # Database connection string credentials (scheme://user:pass@host) r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)", @@ -56,13 +63,21 @@ def _build_secret_patterns() -> re.Pattern: # Catches secrets inside dicts/config dumps by matching on the KEY name # regardless of what the value looks like. # e.g. 'master_key': 'any-value-here', "database_url": "postgres://..." + # private_key with PEM-aware value capture + r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""", r"(?:master_key|database_url|db_url|connection_string|" - r"private_key|signing_key|encryption_key|" + r"signing_key|encryption_key|" r"auth_token|access_token|refresh_token|" r"slack_webhook_url|webhook_url|" r"database_connection_string|" r"huggingface_token|jwt_secret)" r"""['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+""", + # ── Raw JWTs (without Bearer prefix) ── + r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", + # ── Azure SAS tokens in URLs ── + r"[?&]sig=[A-Za-z0-9%+/=]+", + # ── Full JSON service-account blobs (single-line and multi-line) ── + r'\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', ] return re.compile("|".join(patterns), re.IGNORECASE) @@ -74,6 +89,23 @@ def _redact_string(value: str) -> str: return _SECRET_RE.sub(_REDACTED, value) +def redact_secrets(value: str) -> str: + """Public API: redact known secret/credential patterns from an arbitrary string. + + Use this for code paths that bypass the logging system — e.g. Slack/Teams + alerting, HTTP error response bodies, or any other string that may contain + secrets and will be sent to an external sink. + + Not to be confused with redact_message_input_output_from_logging() in + litellm_core_utils/redact_messages.py, which redacts LLM prompt/response + content for privacy — this function redacts credential patterns (API keys, + PEM blocks, tokens, etc.) by shape. + """ + if not _ENABLE_SECRET_REDACTION: + return value + return _redact_string(value) + + class SecretRedactionFilter(logging.Filter): """Scrubs known secret/credential patterns from log records.""" @@ -441,7 +473,7 @@ def _enable_debugging(): def print_verbose(print_statement): try: if set_verbose: - print(print_statement) # noqa + print(redact_secrets(str(print_statement))) # noqa except Exception: pass diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py index d4c5f6a2985..46253bbcf78 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py @@ -2,7 +2,7 @@ Pydantic AI provider configuration. """ -from typing import Any, AsyncIterator, Dict +from typing import Any, AsyncIterator, Dict, Optional from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler @@ -20,10 +20,12 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): self, request_id: str, params: Dict[str, Any], - api_base: str, + api_base: Optional[str] = None, **kwargs, ) -> Dict[str, Any]: """Handle non-streaming request to Pydantic AI agent.""" + if not api_base: + raise ValueError("api_base is required for Pydantic AI agents") return await PydanticAIHandler.handle_non_streaming( request_id=request_id, params=params, @@ -35,10 +37,12 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): self, request_id: str, params: Dict[str, Any], - api_base: str, + api_base: Optional[str] = None, **kwargs, ) -> AsyncIterator[Dict[str, Any]]: """Handle streaming request with fake streaming.""" + if not api_base: + raise ValueError("api_base is required for Pydantic AI agents") async for chunk in PydanticAIHandler.handle_streaming( request_id=request_id, params=params, diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index a638a28aba3..1423617cac0 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -82,6 +82,8 @@ class MCPSigV4Auth(httpx.Auth): aws_session_token: Optional[str] = None, aws_region_name: Optional[str] = None, aws_service_name: Optional[str] = None, + aws_role_name: Optional[str] = None, + aws_session_name: Optional[str] = None, ): try: from botocore.credentials import Credentials @@ -97,7 +99,16 @@ class MCPSigV4Auth(httpx.Auth): # Note: os.environ/ prefixed values are already resolved by # ProxyConfig._check_for_os_environ_vars() at config load time. # Values arrive here as plain strings. - if aws_access_key_id and aws_secret_access_key: + if aws_role_name: + self.credentials = self._assume_role( + aws_role_name=aws_role_name, + aws_session_name=aws_session_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=self.region_name, + ) + elif aws_access_key_id and aws_secret_access_key: self.credentials = Credentials( access_key=aws_access_key_id, secret_key=aws_secret_access_key, @@ -116,6 +127,43 @@ class MCPSigV4Auth(httpx.Auth): "(env vars, ~/.aws/credentials, instance profile)." ) + @staticmethod + def _assume_role( + aws_role_name: str, + aws_session_name: Optional[str], + aws_access_key_id: Optional[str], + aws_secret_access_key: Optional[str], + aws_session_token: Optional[str], + aws_region_name: str, + ): + """Call STS AssumeRole and return temporary credentials.""" + import boto3 + from botocore.credentials import Credentials + + session_name = ( + aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" + ) + + sts_kwargs: dict = {"region_name": aws_region_name} + if aws_access_key_id and aws_secret_access_key: + sts_kwargs["aws_access_key_id"] = aws_access_key_id + sts_kwargs["aws_secret_access_key"] = aws_secret_access_key + if aws_session_token: + sts_kwargs["aws_session_token"] = aws_session_token + + sts_client = boto3.client("sts", **sts_kwargs) + sts_response = sts_client.assume_role( + RoleArn=aws_role_name, + RoleSessionName=session_name, + ) + + sts_creds = sts_response["Credentials"] + return Credentials( + access_key=sts_creds["AccessKeyId"], + secret_key=sts_creds["SecretAccessKey"], + token=sts_creds["SessionToken"], + ) + def auth_flow( self, request: httpx.Request ) -> Generator[httpx.Request, httpx.Response, None]: diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 6fc7b9c1048..50c1cd9d989 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -275,12 +275,11 @@ class AzureBlobStorageLogger(CustomBatchLogger): """ Gets Azure AD token to use for Azure Storage API requests """ - verbose_logger.debug("Getting Azure AD Token from Azure Storage") verbose_logger.debug( - "tenant_id %s, client_id %s, client_secret %s", + "Getting Azure AD Token from Azure Storage, tenant_id=%s, client_id=%s, client_secret=[set=%s]", tenant_id, client_id, - client_secret, + client_secret is not None, ) if tenant_id is None: raise ValueError( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py index 923f613291f..0089e54b1c2 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py @@ -70,7 +70,9 @@ class GCSBucketBase(CustomBatchLogger): custom_llm_provider="vertex_ai", api_base=None, ) - verbose_logger.debug("constructed auth_header %s", auth_header) + verbose_logger.debug( + "constructed auth_header [set=%s]", auth_header is not None + ) headers = { "Authorization": f"Bearer {auth_header}", # auth_header "Content-Type": "application/json", @@ -106,7 +108,9 @@ class GCSBucketBase(CustomBatchLogger): custom_llm_provider="vertex_ai", api_base=None, ) - verbose_logger.debug("constructed auth_header %s", auth_header) + verbose_logger.debug( + "constructed auth_header [set=%s]", auth_header is not None + ) headers = { "Authorization": f"Bearer {auth_header}", # auth_header "Content-Type": "application/json", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 0a2f07bcb25..95bcd4d7186 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -202,8 +202,8 @@ def get_llm_provider( # noqa: PLR0915 ) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key + "dynamic_api_key needs to be a string. Got type={}".format( + type(dynamic_api_key).__name__ ) ) return model, custom_llm_provider, dynamic_api_key, api_base diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index fcdb3eca23a..4fc1ae960b8 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -101,17 +101,15 @@ def get_azure_ad_token_from_entra_id( _client_secret = client_secret verbose_logger.debug( - "tenant_id %s, client_id %s, client_secret %s", + "tenant_id=%s, client_id=%s, client_secret=[set=%s]", _tenant_id, _client_id, - _client_secret, + _client_secret is not None, ) if _tenant_id is None or _client_id is None or _client_secret is None: raise ValueError("tenant_id, client_id, and client_secret must be provided") credential = ClientSecretCredential(_tenant_id, _client_id, _client_secret) - verbose_logger.debug("credential %s", credential) - token_provider = get_bearer_token_provider(credential, scope) verbose_logger.debug("token_provider %s", token_provider) @@ -140,10 +138,10 @@ def get_azure_ad_token_from_username_password( from azure.identity import UsernamePasswordCredential, get_bearer_token_provider verbose_logger.debug( - "client_id %s, azure_username %s, azure_password %s", + "client_id=%s, azure_username=[set=%s], azure_password=[set=%s]", client_id, - azure_username, - azure_password, + azure_username is not None, + azure_password is not None, ) credential = UsernamePasswordCredential( client_id=client_id, @@ -151,8 +149,6 @@ def get_azure_ad_token_from_username_password( password=azure_password, ) - verbose_logger.debug("credential %s", credential) - token_provider = get_bearer_token_provider(credential, scope) verbose_logger.debug("token_provider %s", token_provider) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index b159d62367d..4157fac53b8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -156,24 +156,24 @@ class BaseAWSLLM: verbose_logger.debug( "in get credentials\n" - "aws_access_key_id=%s\n" - "aws_secret_access_key=%s\n" - "aws_session_token=%s\n" + "aws_access_key_id=[set=%s]\n" + "aws_secret_access_key=[set=%s]\n" + "aws_session_token=[set=%s]\n" "aws_region_name=%s\n" "aws_session_name=%s\n" "aws_profile_name=%s\n" "aws_role_name=%s\n" - "aws_web_identity_token=%s\n" + "aws_web_identity_token=[set=%s]\n" "aws_sts_endpoint=%s\n" "aws_external_id=%s", - aws_access_key_id, - aws_secret_access_key, - aws_session_token, + aws_access_key_id is not None, + aws_secret_access_key is not None, + aws_session_token is not None, aws_region_name, aws_session_name, aws_profile_name, aws_role_name, - aws_web_identity_token, + aws_web_identity_token is not None, aws_sts_endpoint, aws_external_id, ) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index b1af7ed2ec3..79cd1c00606 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -174,10 +174,15 @@ def load_private_key_from_file(file_path: str): def get_vendor_from_model(model: str) -> OCIVendors: """ Extracts the vendor from the model name. + + OCI GenAI API uses two apiFormat values: + - "COHERE" for Cohere models (command-r, command-a, etc.) + - "GENERIC" for all other models (Meta Llama, xAI Grok, Google Gemini, etc.) + Args: - model (str): The model name. + model (str): The model name (e.g., "cohere.command-a-03-2025", "meta.llama-3.3-70b-instruct"). Returns: - str: The vendor name. + OCIVendors: The vendor enum value. """ vendor = model.split(".")[0].lower() if vendor == "cohere": diff --git a/litellm/llms/oci/embed/__init__.py b/litellm/llms/oci/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/oci/embed/transformation.py b/litellm/llms/oci/embed/transformation.py new file mode 100644 index 00000000000..1dcd8c5213c --- /dev/null +++ b/litellm/llms/oci/embed/transformation.py @@ -0,0 +1,347 @@ +""" +OCI Generative AI Embedding Configuration + +Supports embedding models available on Oracle Cloud Infrastructure Generative AI service. +Uses the same authentication mechanisms as OCI chat (manual signing or OCI SDK Signer). + +Supported models: +- cohere.embed-english-v3.0 +- cohere.embed-english-light-v3.0 +- cohere.embed-multilingual-v3.0 +- cohere.embed-multilingual-light-v3.0 +- cohere.embed-english-image-v3.0 +- cohere.embed-english-light-image-v3.0 +- cohere.embed-multilingual-light-image-v3.0 +- cohere.embed-v4.0 + +Reference: https://docs.oracle.com/en-us/iaas/api/#/en/generative-ai-inference/latest/EmbedTextResult/EmbedText +""" + +from typing import Any, Dict, List, Optional, Union + +import httpx + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.llms.oci.chat.transformation import OCIChatConfig +from litellm.llms.oci.common_utils import OCIError +from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues +from litellm.types.utils import EmbeddingResponse, Usage + +# Input type mapping from OpenAI conventions to OCI/Cohere conventions +_INPUT_TYPE_MAP = { + "search_document": "SEARCH_DOCUMENT", + "search_query": "SEARCH_QUERY", + "classification": "CLASSIFICATION", + "clustering": "CLUSTERING", +} + + +class OCIEmbeddingConfig(BaseEmbeddingConfig): + """ + Configuration for OCI Generative AI Embedding API. + + The OCI embedding endpoint uses the Cohere embed models hosted on OCI. + Authentication is handled via OCI request signing (manual credentials or OCI SDK Signer). + + Usage: + ```python + import litellm + + response = litellm.embedding( + model="oci/cohere.embed-english-v3.0", + input=["Hello world", "Goodbye world"], + oci_compartment_id="ocid1.compartment.oc1..xxx", + oci_region="us-ashburn-1", + oci_user="ocid1.user.oc1..xxx", + oci_fingerprint="xx:xx:xx:xx", + oci_tenancy="ocid1.tenancy.oc1..xxx", + oci_key_file="~/.oci/key.pem", + ) + ``` + """ + + def __init__(self) -> None: + # We reuse OCIChatConfig for signing logic + self._chat_config = OCIChatConfig() + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + if api_base: + return api_base + + oci_region = optional_params.get("oci_region", "us-ashburn-1") + return f"https://inference.generativeai.{oci_region}.oci.oraclecloud.com/20231130/actions/embedText" + + def get_supported_openai_params(self, model: str) -> list: + return [ + "dimensions", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + # Note: OCI Cohere embed does not support custom dimensions natively, + # but we pass it through in case future models support it + if "dimensions" in non_default_params: + optional_params["dimensions"] = non_default_params["dimensions"] + return optional_params + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate OCI credentials for embedding requests. + Supports both OCI SDK Signer and manual credential signing. + """ + oci_signer = optional_params.get("oci_signer") + oci_region = optional_params.get("oci_region", "us-ashburn-1") + + api_base = ( + api_base + or f"https://inference.generativeai.{oci_region}.oci.oraclecloud.com" + ) + + if oci_signer is None: + oci_user = optional_params.get("oci_user") + oci_fingerprint = optional_params.get("oci_fingerprint") + oci_tenancy = optional_params.get("oci_tenancy") + oci_key = optional_params.get("oci_key") + oci_key_file = optional_params.get("oci_key_file") + oci_compartment_id = optional_params.get("oci_compartment_id") + + if ( + not oci_user + or not oci_fingerprint + or not oci_tenancy + or not (oci_key or oci_key_file) + or not oci_compartment_id + ): + raise Exception( + "Missing required parameters: oci_user, oci_fingerprint, oci_tenancy, oci_compartment_id " + "and at least one of oci_key or oci_key_file. " + "Alternatively, provide an oci_signer object from the OCI SDK." + ) + + from litellm.llms.custom_httpx.http_handler import version + + headers.update( + { + "content-type": "application/json", + "user-agent": f"litellm/{version}", + } + ) + + return headers + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ): + """Delegate to OCIChatConfig's signing logic.""" + return self._chat_config.sign_request( + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + api_base: Optional[str] = None, + ) -> dict: + """ + Transform the embedding request to OCI format. + + OCI embedText API expects: + { + "compartmentId": "...", + "servingMode": {"servingType": "ON_DEMAND", "modelId": "..."}, + "inputs": ["text1", "text2"], + "truncate": "END", + "inputType": "SEARCH_DOCUMENT" + } + """ + oci_compartment_id = optional_params.get("oci_compartment_id") + if not oci_compartment_id: + raise Exception( + "kwarg `oci_compartment_id` is required for OCI embedding requests" + ) + + # Build serving mode + oci_serving_mode = optional_params.get("oci_serving_mode", "ON_DEMAND") + if oci_serving_mode == "DEDICATED": + oci_endpoint_id = optional_params.get("oci_endpoint_id", model) + serving_mode = { + "servingType": "DEDICATED", + "endpointId": oci_endpoint_id, + } + else: + serving_mode = { + "servingType": "ON_DEMAND", + "modelId": model, + } + + # Normalize input to list of strings + if isinstance(input, str): + inputs = [input] + elif isinstance(input, list): + inputs = [] + for item in input: + if isinstance(item, str): + inputs.append(item) + elif isinstance(item, list): + raise ValueError( + "OCI embedding does not support token-array inputs. " + "Please convert token lists to strings before calling embedding()." + ) + else: + inputs.append(str(item)) + else: + inputs = [str(input)] + + # Build request data — OCI embedText API expects inputs, truncate, + # and inputType at the top level alongside compartmentId and servingMode + request_data: Dict[str, Any] = { + "compartmentId": oci_compartment_id, + "servingMode": serving_mode, + "inputs": inputs, + "truncate": optional_params.get("truncate", "END"), + } + + # Map input_type if provided + input_type = optional_params.get("input_type") + if input_type: + mapped_type = _INPUT_TYPE_MAP.get(input_type.lower(), input_type.upper()) + request_data["inputType"] = mapped_type + + # Sign the request using the same URL the HTTP handler will POST to + signing_url = self.get_complete_url( + api_base=api_base, + api_key=None, + model=model, + optional_params=optional_params, + litellm_params={}, + ) + + signed_headers, body = self.sign_request( + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=signing_url, + ) + headers.update(signed_headers) + + return request_data + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> EmbeddingResponse: + """ + Transform OCI embedding response to standard EmbeddingResponse format. + + OCI response format: + { + "embeddings": [[0.1, 0.2, ...], [0.3, 0.4, ...]], + "modelId": "cohere.embed-english-v3.0", + "modelVersion": "3.0", + "inputTextTokenCounts": [5, 4] + } + """ + if raw_response.status_code != 200: + raise OCIError( + message=raw_response.text, + status_code=raw_response.status_code, + ) + + try: + raw_response_json = raw_response.json() + except Exception: + raise OCIError( + message=raw_response.text, + status_code=raw_response.status_code, + ) + + embeddings = raw_response_json.get("embeddings", []) + model_id = raw_response_json.get("modelId", model) + + # Build response data in OpenAI format + embedding_data = [] + for idx, embedding in enumerate(embeddings): + embedding_data.append( + { + "object": "embedding", + "index": idx, + "embedding": embedding, + } + ) + + model_response.model = model_id + model_response.data = embedding_data + model_response.object = "list" + + # Calculate token usage + input_token_counts = raw_response_json.get("inputTextTokenCounts", []) + total_tokens = sum(input_token_counts) if input_token_counts else 0 + + usage = Usage( + prompt_tokens=total_tokens, + total_tokens=total_tokens, + ) + model_response.usage = usage + + return model_response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: Union[dict, httpx.Headers], + ) -> BaseLLMException: + return OCIError( + message=error_message, + status_code=status_code, + headers=headers if isinstance(headers, httpx.Headers) else None, + ) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 1a29ba82eac..cdabac27af7 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -81,26 +81,26 @@ class VertexBase: ) -> Tuple[Any, str]: if credentials is not None: if isinstance(credentials, str): + _is_path = os.path.exists( + credentials + ) # credentials is from server config (litellm_params), not user input verbose_logger.debug( - "Vertex: Loading vertex credentials from %s", credentials - ) - verbose_logger.debug( - "Vertex: checking if credentials is a valid path, os.path.exists(%s)=%s, current dir %s", - credentials, - os.path.exists(credentials), + "Vertex: Loading vertex credentials, is_file_path=%s, current dir %s", + _is_path, os.getcwd(), ) try: - if os.path.exists(credentials): - json_obj = json.load(open(credentials)) + if _is_path: + with open(credentials) as f: + json_obj = json.load(f) else: json_obj = json.loads(credentials) - except Exception: + except Exception as e: raise Exception( - "Unable to load vertex credentials from environment. Got={}".format( - credentials - ) + "Unable to load vertex credentials from environment. " + "Ensure the JSON is valid (check for unescaped newlines in private_key). " + "Parse error: {}".format(type(e).__name__) ) elif isinstance(credentials, dict): json_obj = credentials @@ -668,8 +668,8 @@ class VertexBase: ## VALIDATION STEP if _credentials.token is None or not isinstance(_credentials.token, str): raise ValueError( - "Could not resolve credentials token. Got None or non-string token - {}".format( - _credentials.token + "Could not resolve credentials token. Got None or non-string token (type={})".format( + type(_credentials.token).__name__ ) ) diff --git a/litellm/main.py b/litellm/main.py index eace9c630ba..cbedd1735c7 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5668,6 +5668,22 @@ def embedding( # noqa: PLR0915 aembedding=aembedding, litellm_params={}, ) + elif custom_llm_provider == "oci": + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) elif custom_llm_provider in litellm._custom_providers: custom_handler: Optional[CustomLLM] = None for item in litellm.custom_provider_map: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4e1c7f4ac80..8351f084f28 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23803,7 +23803,8 @@ "output_cost_per_token": 2e-06, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", "supports_function_calling": true, - "supports_response_schema": false + "supports_response_schema": false, + "supports_vision": true }, "oci/meta.llama-3.3-70b-instruct": { "input_cost_per_token": 7.2e-07, @@ -23937,6 +23938,287 @@ "supports_function_calling": true, "supports_response_schema": false }, + "oci/cohere.command-a-reasoning-08-2025": { + "input_cost_per_token": 1.56e-06, + "litellm_provider": "oci", + "max_input_tokens": 256000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.56e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/cohere.command-a-vision-07-2025": { + "input_cost_per_token": 1.56e-06, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.56e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true + }, + "oci/cohere.command-a-translate-08-2025": { + "input_cost_per_token": 9e-08, + "litellm_provider": "oci", + "max_input_tokens": 256000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 9e-08, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": false, + "supports_response_schema": false + }, + "oci/cohere.command-r-08-2024": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/cohere.command-r-plus-08-2024": { + "input_cost_per_token": 1.56e-06, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.56e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/meta.llama-3.2-11b-vision-instruct": { + "input_cost_per_token": 2e-06, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true + }, + "oci/meta.llama-3.1-70b-instruct": { + "input_cost_per_token": 7.2e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/meta.llama-3.3-70b-instruct-fp8-dynamic": { + "input_cost_per_token": 7.2e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4-fast": { + "input_cost_per_token": 5e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4.1-fast": { + "input_cost_per_token": 5e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4.20": { + "input_cost_per_token": 3e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4.20-multi-agent": { + "input_cost_per_token": 3e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-code-fast-1": { + "input_cost_per_token": 5e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/google.gemini-2.5-pro": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "oci", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "oci/google.gemini-2.5-flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "oci", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "oci/google.gemini-2.5-flash-lite": { + "input_cost_per_token": 7.5e-08, + "litellm_provider": "oci", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "oci/cohere.embed-english-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-english-light-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-multilingual-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-multilingual-light-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-english-image-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, + "oci/cohere.embed-english-light-image-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, + "oci/cohere.embed-multilingual-light-image-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, + "oci/cohere.embed-v4.0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, "ollama/codegeex4": { "input_cost_per_token": 0.0, "litellm_provider": "ollama", @@ -30305,6 +30587,27 @@ "supports_pdf_input": true, "supports_tool_choice": true }, + "vertex_ai/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/haiku-4-5", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_native_streaming": true, + "supports_vision": true + }, "vertex_ai/claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, "cache_read_input_token_cost": 1e-07, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1e9d5c5a529..cfbc2c437b0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -346,6 +346,8 @@ class MCPServerManager: aws_session_token=server_config.get("aws_session_token", None), aws_region_name=server_config.get("aws_region_name", None), aws_service_name=server_config.get("aws_service_name", None), + aws_role_name=server_config.get("aws_role_name", None), + aws_session_name=server_config.get("aws_session_name", None), ) self.config_mcp_servers[server_id] = new_server @@ -686,6 +688,8 @@ class MCPServerManager: aws_session_token=aws_creds.get("aws_session_token"), aws_region_name=aws_creds.get("aws_region_name"), aws_service_name=aws_creds.get("aws_service_name"), + aws_role_name=aws_creds.get("aws_role_name"), + aws_session_name=aws_creds.get("aws_session_name"), ) return new_server @@ -1011,6 +1015,8 @@ class MCPServerManager: aws_session_token=server.aws_session_token, aws_region_name=server.aws_region_name, aws_service_name=server.aws_service_name, + aws_role_name=server.aws_role_name, + aws_session_name=server.aws_session_name, ) return MCPClient( @@ -1571,6 +1577,8 @@ class MCPServerManager: ), "aws_region_name": credentials_dict.get("aws_region_name"), "aws_service_name": credentials_dict.get("aws_service_name"), + "aws_role_name": credentials_dict.get("aws_role_name"), + "aws_session_name": credentials_dict.get("aws_session_name"), } def _extract_scopes(self, scopes_value: Any) -> Optional[List[str]]: diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index bb00141ad0a..10b1759b77e 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -136,7 +136,9 @@ class Oauth2Handler: + CommonProxyErrors.not_premium_user.value ) - verbose_proxy_logger.debug("Oauth2 token validation for token=%s", token) + verbose_proxy_logger.debug( + "Oauth2 token validation for token=[set=%s]", token is not None + ) # Get the token info endpoint from environment variable token_info_endpoint = os.getenv("OAUTH_TOKEN_INFO_ENDPOINT") diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index 7e517092b8a..0dc696bc455 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -37,9 +37,13 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: else: auth_data[key] = value verbose_proxy_logger.debug( - f"Auth data before creating UserAPIKeyAuth object: {auth_data}" + "Auth data before creating UserAPIKeyAuth object: keys=%s", + list(auth_data.keys()), ) user_api_key_auth = UserAPIKeyAuth(**auth_data) - verbose_proxy_logger.debug(f"UserAPIKeyAuth object created: {user_api_key_auth}") + verbose_proxy_logger.debug( + "UserAPIKeyAuth object created with keys: %s", + list(user_api_key_auth.__fields_set__), + ) # Create and return UserAPIKeyAuth object return user_api_key_auth diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index fefc6c8af9c..b26d8336191 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -711,7 +711,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "global_max_parallel_requests", None ) user_api_key = _metadata.get("user_api_key", None) - self.print_verbose(f"user_api_key: {user_api_key}") + self.print_verbose(f"user_api_key: [set={user_api_key is not None}]") if user_api_key is None: return diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 497ddcbbe84..00d8ce182ec 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4228,13 +4228,19 @@ async def list_keys( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), page: int = Query(1, description="Page number", ge=1), size: int = Query(10, description="Page size", ge=1, le=100), - user_id: Optional[str] = Query(None, description="Filter keys by user ID. Supports partial matching (substring, case-insensitive)."), + user_id: Optional[str] = Query( + None, + description="Filter keys by user ID. Supports partial matching (substring, case-insensitive).", + ), team_id: Optional[str] = Query(None, description="Filter keys by team ID"), organization_id: Optional[str] = Query( None, description="Filter keys by organization ID" ), key_hash: Optional[str] = Query(None, description="Filter keys by key hash"), - key_alias: Optional[str] = Query(None, description="Filter keys by key alias. Supports partial matching (substring, case-insensitive)."), + key_alias: Optional[str] = Query( + None, + description="Filter keys by key alias. Supports partial matching (substring, case-insensitive).", + ), return_full_object: bool = Query(False, description="Return full key object"), include_team_keys: bool = Query( False, description="Include all keys for teams that user is an admin of." diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8a122a13506..2a661a8348e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -375,9 +375,7 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - user_update, -) +from litellm.proxy.management_endpoints.internal_user_endpoints import user_update from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( router as jwt_key_mapping_router, ) @@ -446,9 +444,7 @@ from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_route from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - set_files_config, -) +from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -552,9 +548,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import ( - DeploymentTypedDict, -) +from litellm.types.router import DeploymentTypedDict from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -9309,20 +9303,17 @@ async def _add_access_group_models_to_team_models( return team_models # Single batch fetch for all access groups - access_group_rows = ( - await prisma_client.db.litellm_accessgrouptable.find_many( - where={"access_group_id": {"in": list(all_access_group_ids)}} - ) + access_group_rows = await prisma_client.db.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": list(all_access_group_ids)}} ) ag_model_map: Dict[str, List[str]] = { - row.access_group_id: row.access_model_names or [] - for row in access_group_rows + row.access_group_id: row.access_model_names or [] for row in access_group_rows } # Second pass: resolve deployments for each eligible team for team_object in eligible_teams: model_names: Set[str] = set() - for ag_id in team_object.access_group_ids or [] : + for ag_id in team_object.access_group_ids or []: model_names.update(ag_model_map.get(ag_id, [])) for model_name in model_names: @@ -9333,9 +9324,7 @@ async def _add_access_group_models_to_team_models( for deployment in deployments: model_id = deployment.get("model_info", {}).get("id", None) if model_id is not None: - team_models.setdefault(model_id, set()).add( - team_object.team_id - ) + team_models.setdefault(model_id, set()).add(team_object.team_id) return team_models diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 97fa6c39fb5..3c1b7cfd10c 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1013,7 +1013,9 @@ async def get_global_spend_report( "/spend/report endpoint " + CommonProxyErrors.not_premium_user.value ) if api_key is not None: - verbose_proxy_logger.debug("Getting /spend for api_key: %s", api_key) + verbose_proxy_logger.debug( + "Getting /spend for api_key: [set=%s]", api_key is not None + ) if api_key.startswith("sk-"): api_key = hash_token(token=api_key) sql_query = """ diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 845919e9120..34c22d2deeb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3293,7 +3293,7 @@ class PrismaClient: if update_key_values is not None: update_key_values = self.jsonify_object(data=update_key_values) if token is not None: - print_verbose(f"token: {token}") + print_verbose(f"token: [set={token is not None}]") # check if plain text or hash token = _hash_token_if_needed(token=token) db_data["token"] = token diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 9075373f1cf..2207acbb37a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2123,9 +2123,9 @@ class LiteLLMCompletionResponsesConfig: hasattr(completion_details, "reasoning_tokens") and completion_details.reasoning_tokens is not None ): - output_details_dict["reasoning_tokens"] = ( - completion_details.reasoning_tokens - ) + output_details_dict[ + "reasoning_tokens" + ] = completion_details.reasoning_tokens else: output_details_dict["reasoning_tokens"] = 0 diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 1309846c102..1188ce9d592 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -147,6 +147,13 @@ async def get_deployments_for_tag( ) return healthy_deployments + # Tag filtering applies only when there is at least one deployment to evaluate. + if isinstance(healthy_deployments, list) and len(healthy_deployments) == 0: + verbose_logger.debug( + "get_deployments_for_tag: empty candidate set; skipping tag filter" + ) + return healthy_deployments + verbose_logger.debug( "request metadata: %s", request_kwargs.get(metadata_variable_name) ) diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 63231923f1a..c23e6ce473a 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, Any, Optional, Union -from litellm._logging import verbose_router_logger +from litellm._logging import redact_secrets, verbose_router_logger from litellm.constants import MAX_EXCEPTION_MESSAGE_LENGTH from litellm.router_utils.cooldown_handlers import ( _async_get_cooldown_deployments_with_debug_info, @@ -57,6 +57,9 @@ async def send_llm_exception_alert( exception_str += litellm_debug_info exception_str += f"\n\n{error_traceback_str[:MAX_EXCEPTION_MESSAGE_LENGTH]}" + # Redact secrets before sending to external service (Slack / MS Teams) + exception_str = redact_secrets(exception_str) + await litellm_router_instance.slack_alerting_logger.send_alert( message=f"LLM API call failed: `{exception_str}`", level="High", diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index eb90dda0e99..0b16f7e10ad 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -117,12 +117,14 @@ def get_secret_from_manager( # noqa: PLR0915 secret_name=secret_name, primary_secret_name=primary_secret_name, ) - print_verbose(f"get_secret_value_response: {secret}") + print_verbose(f"get_secret_value_response: [set={secret is not None}]") elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: secret = client.get_secret_from_google_secret_manager(secret_name) - print_verbose(f"secret from google secret manager: {secret}") + print_verbose( + f"secret from google secret manager: [set={secret is not None}]" + ) if secret is None: raise ValueError( f"No secret found in Google Secret Manager for {secret_name}" diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index af91926de2f..ebabf3fb6f8 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -111,6 +111,12 @@ class MCPCredentials(TypedDict, total=False): aws_service_name: Optional[str] """AWS service name for SigV4 signing (e.g., 'bedrock-agentcore'). Not a secret — stored unencrypted.""" + aws_role_name: Optional[str] + """IAM role ARN for STS AssumeRole (e.g., 'arn:aws:iam::123456789012:role/MyRole'). Not a secret — stored unencrypted.""" + + aws_session_name: Optional[str] + """Session name for STS AssumeRole (used in CloudTrail). Not a secret — stored unencrypted.""" + class MCPServerCostInfo(TypedDict, total=False): default_cost_per_query: Optional[float] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index ed391f0af68..db7657a0174 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -54,6 +54,8 @@ class MCPServer(BaseModel): aws_session_token: Optional[str] = None aws_region_name: Optional[str] = None aws_service_name: Optional[str] = None # defaults to "bedrock-agentcore" + aws_role_name: Optional[str] = None # IAM role ARN for STS AssumeRole + aws_session_name: Optional[str] = None # session name for CloudTrail auditing # Stdio-specific fields command: Optional[str] = None args: Optional[List[str]] = None diff --git a/litellm/utils.py b/litellm/utils.py index 37bc35af299..6806961bf51 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8314,6 +8314,10 @@ class ProviderConfigManager: return SagemakerEmbeddingConfig.get_model_config(model) elif litellm.LlmProviders.PERPLEXITY == provider: return litellm.PerplexityEmbeddingConfig() + elif litellm.LlmProviders.OCI == provider: + from litellm.llms.oci.embed.transformation import OCIEmbeddingConfig + + return OCIEmbeddingConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 35fa2206761..6da9a004b8a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23788,7 +23788,8 @@ "output_cost_per_token": 2e-06, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", "supports_function_calling": true, - "supports_response_schema": false + "supports_response_schema": false, + "supports_vision": true }, "oci/meta.llama-3.3-70b-instruct": { "input_cost_per_token": 7.2e-07, @@ -23922,6 +23923,287 @@ "supports_function_calling": true, "supports_response_schema": false }, + "oci/cohere.command-a-reasoning-08-2025": { + "input_cost_per_token": 1.56e-06, + "litellm_provider": "oci", + "max_input_tokens": 256000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.56e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/cohere.command-a-vision-07-2025": { + "input_cost_per_token": 1.56e-06, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.56e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true + }, + "oci/cohere.command-a-translate-08-2025": { + "input_cost_per_token": 9e-08, + "litellm_provider": "oci", + "max_input_tokens": 256000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 9e-08, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": false, + "supports_response_schema": false + }, + "oci/cohere.command-r-08-2024": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/cohere.command-r-plus-08-2024": { + "input_cost_per_token": 1.56e-06, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 1.56e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/meta.llama-3.2-11b-vision-instruct": { + "input_cost_per_token": 2e-06, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true + }, + "oci/meta.llama-3.1-70b-instruct": { + "input_cost_per_token": 7.2e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/meta.llama-3.3-70b-instruct-fp8-dynamic": { + "input_cost_per_token": 7.2e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4-fast": { + "input_cost_per_token": 5e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4.1-fast": { + "input_cost_per_token": 5e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4.20": { + "input_cost_per_token": 3e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-4.20-multi-agent": { + "input_cost_per_token": 3e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/xai.grok-code-fast-1": { + "input_cost_per_token": 5e-06, + "litellm_provider": "oci", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": false + }, + "oci/google.gemini-2.5-pro": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "oci", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "oci/google.gemini-2.5-flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "oci", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "oci/google.gemini-2.5-flash-lite": { + "input_cost_per_token": 7.5e-08, + "litellm_provider": "oci", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "oci/cohere.embed-english-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-english-light-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-multilingual-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-multilingual-light-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing" + }, + "oci/cohere.embed-english-image-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, + "oci/cohere.embed-english-light-image-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, + "oci/cohere.embed-multilingual-light-image-v3.0": { + "input_cost_per_token": 1e-07, + "litellm_provider": "oci", + "max_input_tokens": 512, + "max_tokens": 512, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 384, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, + "oci/cohere.embed-v4.0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "oci", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", + "supports_embedding_image_input": true + }, "ollama/codegeex4": { "input_cost_per_token": 0.0, "litellm_provider": "ollama", @@ -30290,6 +30572,27 @@ "supports_pdf_input": true, "supports_tool_choice": true }, + "vertex_ai/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/haiku-4-5", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_native_streaming": true, + "supports_vision": true + }, "vertex_ai/claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, "cache_read_input_token_cost": 1e-07, diff --git a/tests/test_litellm/llms/oci/embed/__init__.py b/tests/test_litellm/llms/oci/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/oci/embed/test_oci_embedding.py b/tests/test_litellm/llms/oci/embed/test_oci_embedding.py new file mode 100644 index 00000000000..4ecca377e63 --- /dev/null +++ b/tests/test_litellm/llms/oci/embed/test_oci_embedding.py @@ -0,0 +1,369 @@ +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.oci.embed.transformation import OCIEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +# Test constants +TEST_MODEL_NAME = "cohere.embed-english-v3.0" +TEST_MODEL = f"oci/{TEST_MODEL_NAME}" +TEST_COMPARTMENT_ID = "ocid1.compartment.oc1..xxxxxx" +BASE_OCI_PARAMS = { + "oci_region": "us-ashburn-1", + "oci_user": "ocid1.user.oc1..xxxxxxEXAMPLExxxxxx", + "oci_fingerprint": "4f:29:77:cc:b1:3e:55:ab:61:2a:de:47:f1:38:4c:90", + "oci_tenancy": "ocid1.tenancy.oc1..xxxxxxEXAMPLExxxxxx", + "oci_compartment_id": TEST_COMPARTMENT_ID, +} + +TEST_OCI_PARAMS_KEY = { + **BASE_OCI_PARAMS, + "oci_key": "", +} + +TEST_OCI_PARAMS_KEY_FILE = { + **BASE_OCI_PARAMS, + "oci_key_file": "", +} + +# Mock OCI embedding response +MOCK_OCI_EMBEDDING_RESPONSE = { + "embeddings": [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]], + "modelId": "cohere.embed-english-v3.0", + "modelVersion": "3.0", + "inputTextTokenCounts": [5, 4], +} + + +@pytest.fixture(params=[TEST_OCI_PARAMS_KEY, TEST_OCI_PARAMS_KEY_FILE]) +def supplied_params(request): + """Fixture for passing in optional_parameters""" + return request.param + + +class TestOCIEmbeddingConfig: + def test_get_complete_url_default_region(self): + """test_get_complete_url returns URL with us-ashburn-1 when no api_base is given.""" + config = OCIEmbeddingConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model=TEST_MODEL_NAME, + optional_params={}, + litellm_params={}, + ) + assert "us-ashburn-1" in url + assert "embedText" in url + + def test_get_complete_url_custom_region(self): + """test_get_complete_url uses region from optional_params.""" + config = OCIEmbeddingConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model=TEST_MODEL_NAME, + optional_params={"oci_region": "us-chicago-1"}, + litellm_params={}, + ) + assert "us-chicago-1" in url + assert "embedText" in url + + def test_get_complete_url_custom_api_base(self): + """test_get_complete_url returns api_base as-is when provided.""" + config = OCIEmbeddingConfig() + custom_base = "https://custom.oci.example.com/embed" + url = config.get_complete_url( + api_base=custom_base, + api_key=None, + model=TEST_MODEL_NAME, + optional_params={}, + litellm_params={}, + ) + assert url == custom_base + + def test_get_supported_openai_params(self): + """test_get_supported_openai_params returns expected params list.""" + config = OCIEmbeddingConfig() + params = config.get_supported_openai_params(model=TEST_MODEL_NAME) + assert "dimensions" in params + assert "encoding_format" not in params + + def test_map_openai_params_dimensions(self): + """test dimensions is mapped correctly.""" + config = OCIEmbeddingConfig() + optional_params = {} + result = config.map_openai_params( + non_default_params={"dimensions": 512}, + optional_params=optional_params, + model=TEST_MODEL_NAME, + drop_params=False, + ) + assert result["dimensions"] == 512 + + def test_validate_environment_with_credentials(self, supplied_params): + """test validate_environment returns content-type and user-agent headers when credentials are supplied.""" + config = OCIEmbeddingConfig() + headers = {} + result = config.validate_environment( + headers=headers, + model=TEST_MODEL, + messages=[], + optional_params=supplied_params, + litellm_params={}, + ) + assert result["content-type"] == "application/json" + assert "litellm" in result["user-agent"] + + def test_validate_environment_missing_credentials(self): + """test validate_environment raises Exception with 'Missing required parameters' when credentials are incomplete.""" + config = OCIEmbeddingConfig() + incomplete_params = { + "oci_user": "ocid1.user.oc1..xxx", + # Missing oci_fingerprint, oci_tenancy, oci_key/oci_key_file, oci_compartment_id + } + with pytest.raises(Exception) as excinfo: + config.validate_environment( + headers={}, + model=TEST_MODEL, + messages=[], + optional_params=incomplete_params, + litellm_params={}, + ) + assert "Missing required parameters" in str(excinfo.value) + + def test_validate_environment_with_signer(self): + """test validate_environment passes when oci_signer is provided.""" + config = OCIEmbeddingConfig() + + class MockSigner: + def do_request_sign(self, request, enforce_content_headers=True): + request.headers["authorization"] = 'Signature version="1"' + + optional_params = { + "oci_signer": MockSigner(), + "oci_region": "us-ashburn-1", + } + result = config.validate_environment( + headers={}, + model=TEST_MODEL, + messages=[], + optional_params=optional_params, + litellm_params={}, + ) + assert result["content-type"] == "application/json" + + def test_transform_embedding_request_on_demand(self): + """test transform_embedding_request builds correct ON_DEMAND OCI request body.""" + config = OCIEmbeddingConfig() + optional_params = { + "oci_compartment_id": TEST_COMPARTMENT_ID, + } + with patch.object(config, "sign_request", return_value=({}, "{}")): + result = config.transform_embedding_request( + model=TEST_MODEL_NAME, + input=["Hello world", "Goodbye world"], + optional_params=optional_params, + headers={}, + ) + + assert result["compartmentId"] == TEST_COMPARTMENT_ID + assert result["servingMode"]["servingType"] == "ON_DEMAND" + assert result["servingMode"]["modelId"] == TEST_MODEL_NAME + assert result["inputs"] == ["Hello world", "Goodbye world"] + assert result["truncate"] == "END" + + def test_transform_embedding_request_dedicated(self): + """test transform_embedding_request builds DEDICATED servingMode with endpointId.""" + config = OCIEmbeddingConfig() + test_endpoint_id = "ocid1.generativeaiendpoint.oc1.us-chicago-1.xxxxxx" + optional_params = { + "oci_compartment_id": TEST_COMPARTMENT_ID, + "oci_serving_mode": "DEDICATED", + "oci_endpoint_id": test_endpoint_id, + } + with patch.object(config, "sign_request", return_value=({}, "{}")): + result = config.transform_embedding_request( + model=TEST_MODEL_NAME, + input=["Hello world"], + optional_params=optional_params, + headers={}, + ) + + assert result["servingMode"]["servingType"] == "DEDICATED" + assert result["servingMode"]["endpointId"] == test_endpoint_id + + def test_transform_embedding_request_input_type(self): + """test input_type=search_query is mapped to SEARCH_QUERY in request data.""" + config = OCIEmbeddingConfig() + optional_params = { + "oci_compartment_id": TEST_COMPARTMENT_ID, + "input_type": "search_query", + } + with patch.object(config, "sign_request", return_value=({}, "{}")): + result = config.transform_embedding_request( + model=TEST_MODEL_NAME, + input=["What is the capital of Brazil?"], + optional_params=optional_params, + headers={}, + ) + + assert result["inputType"] == "SEARCH_QUERY" + + def test_transform_embedding_request_string_input(self): + """test single string input is wrapped in a list.""" + config = OCIEmbeddingConfig() + optional_params = { + "oci_compartment_id": TEST_COMPARTMENT_ID, + } + with patch.object(config, "sign_request", return_value=({}, "{}")): + result = config.transform_embedding_request( + model=TEST_MODEL_NAME, + input="Hello world", + optional_params=optional_params, + headers={}, + ) + + assert isinstance(result["inputs"], list) + assert result["inputs"] == ["Hello world"] + + def test_transform_embedding_request_token_list_raises(self): + """test token-array inputs raise ValueError instead of silent conversion.""" + config = OCIEmbeddingConfig() + optional_params = { + "oci_compartment_id": TEST_COMPARTMENT_ID, + } + with patch.object(config, "sign_request", return_value=({}, "{}")): + with pytest.raises(ValueError, match="does not support token-array"): + config.transform_embedding_request( + model=TEST_MODEL_NAME, + input=[[1234, 5678]], + optional_params=optional_params, + headers={}, + ) + + def test_transform_embedding_response(self): + """test OCI embedding response is correctly transformed into EmbeddingResponse.""" + config = OCIEmbeddingConfig() + mock_response = httpx.Response( + status_code=200, + json=MOCK_OCI_EMBEDDING_RESPONSE, + request=httpx.Request("POST", "https://test.com"), + ) + mock_logging = MagicMock() + model_response = EmbeddingResponse() + + result = config.transform_embedding_response( + model=TEST_MODEL_NAME, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + assert isinstance(result, EmbeddingResponse) + assert result.model == "cohere.embed-english-v3.0" + assert len(result.data) == 2 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3, 0.4] + assert result.data[1]["embedding"] == [0.5, 0.6, 0.7, 0.8] + assert result.data[0]["index"] == 0 + assert result.data[1]["index"] == 1 + # Total tokens: 5 + 4 = 9 + assert result.usage.prompt_tokens == 9 + assert result.usage.total_tokens == 9 + + def test_transform_embedding_response_error(self): + """test non-200 status code raises OCIError.""" + from litellm.llms.oci.common_utils import OCIError + + config = OCIEmbeddingConfig() + mock_response = httpx.Response( + status_code=400, + text="Bad Request", + request=httpx.Request("POST", "https://test.com"), + ) + mock_logging = MagicMock() + model_response = EmbeddingResponse() + + with pytest.raises(OCIError): + config.transform_embedding_response( + model=TEST_MODEL_NAME, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + def test_model_prices_embedding_models(self): + """test all 8 OCI embedding models exist in model_prices_and_context_window.json with mode=embedding.""" + model_prices_path = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "..", + "..", + "model_prices_and_context_window.json", + ) + with open(model_prices_path) as f: + model_prices = json.load(f) + + expected_embedding_models = [ + "oci/cohere.embed-english-v3.0", + "oci/cohere.embed-english-light-v3.0", + "oci/cohere.embed-multilingual-v3.0", + "oci/cohere.embed-multilingual-light-v3.0", + "oci/cohere.embed-english-image-v3.0", + "oci/cohere.embed-english-light-image-v3.0", + "oci/cohere.embed-multilingual-light-image-v3.0", + "oci/cohere.embed-v4.0", + ] + + for model_key in expected_embedding_models: + assert model_key in model_prices, f"Missing model: {model_key}" + assert ( + model_prices[model_key].get("mode") == "embedding" + ), f"Model {model_key} does not have mode='embedding'" + + def test_model_prices_new_chat_models(self): + """test the 16 new OCI chat models exist in model_prices_and_context_window.json with mode=chat.""" + model_prices_path = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "..", + "..", + "model_prices_and_context_window.json", + ) + with open(model_prices_path) as f: + model_prices = json.load(f) + + expected_chat_models = [ + "oci/xai.grok-3", + "oci/xai.grok-3-fast", + "oci/xai.grok-3-mini", + "oci/xai.grok-3-mini-fast", + "oci/xai.grok-4", + "oci/xai.grok-4-fast", + "oci/xai.grok-4.1-fast", + "oci/xai.grok-4.20", + "oci/xai.grok-4.20-multi-agent", + "oci/xai.grok-code-fast-1", + "oci/cohere.command-a-03-2025", + "oci/cohere.command-a-reasoning-08-2025", + "oci/cohere.command-a-vision-07-2025", + "oci/cohere.command-a-translate-08-2025", + "oci/google.gemini-2.5-pro", + "oci/google.gemini-2.5-flash", + ] + + for model_key in expected_chat_models: + assert model_key in model_prices, f"Missing model: {model_key}" + assert ( + model_prices[model_key].get("mode") == "chat" + ), f"Model {model_key} does not have mode='chat'" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index a2295e1271e..7c142e3a771 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -162,6 +162,140 @@ class TestMCPSigV4Auth: assert "x-amz-security-token" in signed_request.headers +class TestMCPSigV4AssumeRole: + """Tests for STS AssumeRole credential resolution in MCPSigV4Auth.""" + + def test_assume_role_with_ambient_credentials(self): + """MCPSigV4Auth calls STS AssumeRole when aws_role_name is provided (no explicit keys).""" + mock_sts = MagicMock() + mock_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASSUMED_KEY", + "SecretAccessKey": "ASSUMED_SECRET", + "SessionToken": "ASSUMED_TOKEN", + "Expiration": "2026-03-30T12:00:00Z", + } + } + + with patch("boto3.client", return_value=mock_sts) as mock_boto3: + auth = MCPSigV4Auth( + aws_role_name="arn:aws:iam::123456789012:role/TestRole", + aws_region_name="us-east-1", + ) + + mock_boto3.assert_called_once_with("sts", region_name="us-east-1") + mock_sts.assume_role.assert_called_once() + call_kwargs = mock_sts.assume_role.call_args[1] + assert call_kwargs["RoleArn"] == "arn:aws:iam::123456789012:role/TestRole" + assert call_kwargs["RoleSessionName"].startswith("litellm-mcp-") + assert auth.credentials.access_key == "ASSUMED_KEY" + assert auth.credentials.secret_key == "ASSUMED_SECRET" + assert auth.credentials.token == "ASSUMED_TOKEN" + + def test_assume_role_with_explicit_source_credentials(self): + """When aws_role_name + explicit keys are provided, keys are used as STS source identity.""" + mock_sts = MagicMock() + mock_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASSUMED_KEY", + "SecretAccessKey": "ASSUMED_SECRET", + "SessionToken": "ASSUMED_TOKEN", + "Expiration": "2026-03-30T12:00:00Z", + } + } + + with patch("boto3.client", return_value=mock_sts) as mock_boto3: + auth = MCPSigV4Auth( + aws_role_name="arn:aws:iam::123456789012:role/TestRole", + aws_access_key_id="SOURCE_KEY", + aws_secret_access_key="SOURCE_SECRET", + aws_region_name="us-west-2", + ) + + mock_boto3.assert_called_once_with( + "sts", + region_name="us-west-2", + aws_access_key_id="SOURCE_KEY", + aws_secret_access_key="SOURCE_SECRET", + ) + assert auth.credentials.access_key == "ASSUMED_KEY" + + def test_assume_role_with_custom_session_name(self): + """Custom aws_session_name is used in the AssumeRole call.""" + mock_sts = MagicMock() + mock_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASSUMED_KEY", + "SecretAccessKey": "ASSUMED_SECRET", + "SessionToken": "ASSUMED_TOKEN", + "Expiration": "2026-03-30T12:00:00Z", + } + } + + with patch("boto3.client", return_value=mock_sts): + MCPSigV4Auth( + aws_role_name="arn:aws:iam::123456789012:role/TestRole", + aws_session_name="regeneron-litellm-prod", + ) + + call_kwargs = mock_sts.assume_role.call_args[1] + assert call_kwargs["RoleSessionName"] == "regeneron-litellm-prod" + + def test_assume_role_signing_works(self): + """Requests are signed correctly with STS-derived credentials.""" + mock_sts = MagicMock() + mock_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "AKIAIOSFODNN7EXAMPLE", + "SecretAccessKey": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "SessionToken": "STS_SESSION_TOKEN", + "Expiration": "2026-03-30T12:00:00Z", + } + } + + with patch("boto3.client", return_value=mock_sts): + auth = MCPSigV4Auth( + aws_role_name="arn:aws:iam::123456789012:role/TestRole", + aws_region_name="us-east-1", + aws_service_name="bedrock-agentcore", + ) + + request = httpx.Request( + method="POST", + url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", + headers={"Content-Type": "application/json"}, + content=b'{"jsonrpc":"2.0","method":"tools/list","id":1}', + ) + + signed_request = next(auth.auth_flow(request)) + assert "Authorization" in signed_request.headers + assert "AWS4-HMAC-SHA256" in signed_request.headers["Authorization"] + assert "x-amz-security-token" in signed_request.headers + + def test_assume_role_takes_precedence_over_explicit_keys(self): + """When both aws_role_name and explicit keys are provided, AssumeRole is used (keys become source identity).""" + mock_sts = MagicMock() + mock_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASSUMED_KEY", + "SecretAccessKey": "ASSUMED_SECRET", + "SessionToken": "ASSUMED_TOKEN", + "Expiration": "2026-03-30T12:00:00Z", + } + } + + with patch("boto3.client", return_value=mock_sts): + auth = MCPSigV4Auth( + aws_role_name="arn:aws:iam::123456789012:role/TestRole", + aws_access_key_id="EXPLICIT_KEY", + aws_secret_access_key="EXPLICIT_SECRET", + ) + + # Credentials should be from AssumeRole, not the explicit keys + assert auth.credentials.access_key == "ASSUMED_KEY" + assert auth.credentials.secret_key == "ASSUMED_SECRET" + + class TestMCPClientSigV4Integration: """Tests for MCPClient with SigV4 auth wired through.""" @@ -319,6 +453,86 @@ class TestMCPServerManagerSigV4: assert client._aws_auth is None + @pytest.mark.asyncio + async def test_load_config_with_aws_role_name(self): + """Config loading correctly parses aws_role_name and aws_session_name.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + config = { + "agentcore_tools": { + "url": "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", + "transport": "http", + "auth_type": "aws_sigv4", + "aws_role_name": "arn:aws:iam::123456789012:role/TestRole", + "aws_session_name": "litellm-prod", + "aws_region_name": "us-east-1", + } + } + + manager = MCPServerManager() + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.aws_role_name == "arn:aws:iam::123456789012:role/TestRole" + assert server.aws_session_name == "litellm-prod" + + @pytest.mark.asyncio + async def test_create_mcp_client_with_role_assumption(self): + """_create_mcp_client passes aws_role_name to MCPSigV4Auth.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + mock_sts = MagicMock() + mock_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASSUMED_KEY", + "SecretAccessKey": "ASSUMED_SECRET", + "SessionToken": "ASSUMED_TOKEN", + "Expiration": "2026-03-30T12:00:00Z", + } + } + + server = MCPServer( + server_id="test-sigv4-role", + name="test_sigv4_role", + server_name="test_sigv4_role", + url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", + transport=MCPTransport.http, + auth_type=MCPAuth.aws_sigv4, + aws_role_name="arn:aws:iam::123456789012:role/TestRole", + aws_region_name="us-east-1", + ) + + manager = MCPServerManager() + with patch("boto3.client", return_value=mock_sts): + client = await manager._create_mcp_client(server=server) + + assert client._aws_auth is not None + assert isinstance(client._aws_auth, MCPSigV4Auth) + mock_sts.assume_role.assert_called_once() + + def test_extract_aws_credentials_includes_role_fields(self): + """_extract_aws_credentials extracts aws_role_name and aws_session_name.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + creds = { + "aws_access_key_id": "KEY", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::123456789012:role/TestRole", + "aws_session_name": "my-session", + } + + result = manager._extract_aws_credentials(creds, credentials_are_encrypted=False) + assert result["aws_role_name"] == "arn:aws:iam::123456789012:role/TestRole" + assert result["aws_session_name"] == "my-session" + class TestSigV4CredentialEncryption: """Test encrypt/decrypt round-trip for AWS SigV4 credentials.""" diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py index ad72a0d1195..f2468d95820 100644 --- a/tests/test_litellm/proxy/db/test_check_migration.py +++ b/tests/test_litellm/proxy/db/test_check_migration.py @@ -27,12 +27,14 @@ def test_check_migration_out_of_sync(mocker): - 🚨 [IMPORTANT] Does NOT Raise an Exception when the Prisma schema is out of sync with the database. - logs an error when the Prisma schema is out of sync with the database. """ - # Mock the logger BEFORE importing the function - mock_logger = mocker.patch("litellm._logging.verbose_logger") - - # Import the function after mocking the logger + # Import the function first so check_migration module is in sys.modules, + # then patch the logger reference in that module directly (not the source + # module) so the patch works regardless of import order or xdist worker + # assignment. from litellm.proxy.db.check_migration import check_prisma_schema_diff + mock_logger = mocker.patch("litellm.proxy.db.check_migration.verbose_logger") + # Mock the helper function to simulate out-of-sync state mock_diff_helper = mocker.patch( "litellm.proxy.db.check_migration.check_prisma_schema_diff_helper", diff --git a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py b/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py index 6c7cfa61b58..dca2bd84f92 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py @@ -188,6 +188,26 @@ async def test_tag_filtering_disabled_returns_all_deployments(): assert result == ALL_DEPLOYMENTS +@pytest.mark.asyncio +async def test_empty_healthy_deployments_with_request_tags_returns_empty_list(): + """ + With an empty candidate list, return [] even when the request includes metadata tags. + + Tag-based filtering runs only against non-empty healthy_deployments; an empty list is + returned unchanged for the router's standard handling. + """ + router = _make_router_mock() + result = await get_deployments_for_tag( + llm_router_instance=router, + model="gpt-5.2", + healthy_deployments=[], + request_kwargs={ + "metadata": {"tags": ["client_id:xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"]} + }, + ) + assert result == [] + + @pytest.mark.asyncio async def test_explicit_tag_match_takes_precedence_over_regex(): """A deployment with both tags and tag_regex: exact tag match fires first.""" diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 1d575195194..6cbb2fd7b8f 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -262,3 +262,80 @@ def test_key_name_redaction_in_general_settings_dict(): assert "REDACTED" in output # Non-sensitive values should survive assert "enable_jwt_auth" in output + + +# ── GCP service-account / Vertex credential redaction ── + + +_SAMPLE_SA_JSON = ( + '{"type": "service_account", "project_id": "my-proj-123", ' + '"private_key_id": "abc123def", ' + '"private_key": "-----BEGIN PRIVATE KEY-----\\nMIIEvQIBADANBgkq\\n-----END PRIVATE KEY-----\\n", ' + '"client_email": "sa@my-proj.iam.gserviceaccount.com", ' + '"client_id": "123456789"}' +) + + +def test_pem_private_key_redacted_in_json(): + result = _redact_string(_SAMPLE_SA_JSON) + assert "MIIEvQIBADA" not in result + assert "-----BEGIN" not in result + + +def test_pem_private_key_redacted_in_dict_repr(): + import json + + sa = json.loads(_SAMPLE_SA_JSON) + result = _redact_string(str(sa)) + assert "MIIEvQIBADA" not in result + + +def test_service_account_blob_fully_redacted(): + result = _redact_string(f"Got={_SAMPLE_SA_JSON}") + assert "my-proj-123" not in result + assert "sa@my-proj.iam.gserviceaccount.com" not in result + assert "abc123def" not in result + assert "MIIEvQIBADA" not in result + + +def test_vertex_error_message_no_credential_leak(): + """The old Vertex error format leaked the full credential JSON. + The new format must not contain any credential material.""" + new_msg = ( + "Unable to load vertex credentials from environment. " + "Ensure the JSON is valid (check for unescaped newlines in private_key). " + "Parse error: JSONDecodeError" + ) + result = _redact_string(new_msg) + assert result == new_msg # nothing to redact + + +def test_vertex_traceback_redacts_pem(): + traceback_text = ( + "Traceback (most recent call last):\n" + ' File "vertex_llm_base.py", line 95\n' + " json_obj = json.loads(credentials)\n" + "json.decoder.JSONDecodeError: Invalid control character\n" + "Failed to load vertex credentials. Error: " + "Unable to load vertex credentials from environment. " + f"Got={_SAMPLE_SA_JSON}" + ) + result = _redact_string(traceback_text) + assert "MIIEvQIBADA" not in result + assert "-----BEGIN" not in result + + +def test_gcp_oauth_token_redacted(): + result = _redact_string("access token ya29.c.c0ASRK0GZvXlongtokenhere") + assert "ya29." not in result + assert "REDACTED" in result + + +def test_non_pem_private_key_value_redacted(): + result = _redact_string("'private_key': 'some-non-pem-secret-value'") + assert "some-non-pem-secret" not in result + + +def test_normal_vertex_log_not_redacted(): + msg = "Vertex: Loading vertex credentials, is_file_path=True, current dir /app" + assert _redact_string(msg) == msg diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index aaddf53fd7e..4c824fcee0b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -934,6 +934,38 @@ const CreateMCPServer: React.FC = ({ className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" /> + + AWS Role ARN + + + + + } + name={["credentials", "aws_role_name"]} + > + + + + AWS Session Name + + + + + } + name={["credentials", "aws_session_name"]} + > + + )} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index 118461f0043..e81d2f3960e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -977,6 +977,38 @@ const MCPServerEdit: React.FC = ({ className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" /> + + AWS Role ARN + + + + + } + name={["credentials", "aws_role_name"]} + > + + + + AWS Session Name + + + + + } + name={["credentials", "aws_session_name"]} + > + + )} From 8d9e5ff3d4f1606dd527aae14f9427a155229e4c Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Sat, 4 Apr 2026 15:13:54 -0700 Subject: [PATCH 40/55] Litellm team model group name routing fix (#25148) (#25154) * Litellm team model group name routing fix (#25148) * fix(team-routing): use deterministic team model group names Use a deterministic internal model_name for team-scoped deployments so sibling deployments with the same public model share a routing group. This makes team alias writes idempotent and preserves multi-deployment failover/load balancing behavior. Made-with: Cursor * fix(team-routing): keep team model routing on public names Remove team model_alias rewrites and resolve team deployments by team_public_model_name with team_id so sibling deployments stay in the routing candidate pool, with explicit logs showing candidate selection before load balancing. Made-with: Cursor * chore(team-routing): remove temporary candidate pool logs Remove temporary fire-emoji router logs used for local verification while keeping team sibling deployment routing behavior unchanged. Made-with: Cursor * fix(router): address Greptile review comments - Add None guard for original_model_name in _add_team_model_to_db - Remove stale old public name when renaming team model - Add comment clarifying team deployment early-return priority Made-with: Cursor * fix(router): address remaining Greptile P0/P1 issues - Update map_team_model test to expect public name return - Only remove old public name if no sibling deployments use it Made-with: Cursor * fix(router): address Greptile P1/P2 performance issues - Guard against llm_router=None to prevent silent deletion - Add O(1) team_model index to avoid O(n) scan on every team request Made-with: Cursor * fix(router): prevent cross-team deployment leakage in fallback path Guard should_include_deployment fallback to only return deployments matching the requested team_id, preventing public-name collisions from leaking deployments across teams Made-with: Cursor * fix(management): query DB directly for sibling deployments on rename - Add clarifying comments to test assertions - Query prisma DB instead of in-memory router to avoid stale state - Prevents incorrect deletion of old public name when siblings exist Made-with: Cursor * fix(router): guard None model_info and deduplicate team index logic - Guard against None model_info in sibling deployment check - Extract _update_team_model_index helper to eliminate duplication Made-with: Cursor * fix(routing): prevent stale model_aliases from interfering with team routing - Skip model_aliases rewrite if model resolves to team deployments - Add test coverage for sibling-preservation branch - Update MockPrismaClient to support sibling deployment scenarios Made-with: Cursor * perf(routing): optimize team model checks and improve test coverage - Use O(1) team index lookup instead of map_team_model in alias guard - Fix MockPrismaClient to validate where clause filters - Add comment explaining DB query trade-off for team deployments Made-with: Cursor * fix(routing): address state consistency and type safety issues - Check alias target pattern to detect stale team aliases - Fix PrismaClient type annotation to Optional - Eliminate in-place mutation in index update logic Made-with: Cursor * Fix greptile comments * Fix greptile comments * Fix greptile comments * Fix greptile comments * Fix greptile comments * Fix code qa issues * Fix greptile reviews and mock test * Fix greptile reviews and mock test * Fix greptile reviews and mock test * fix(router): address Greptile P1/P2 review comments - Add deduplication guard in _update_team_model_index to prevent duplicate indices - Add wildcard comment in map_team_model for clarity - Add monkeypatch to test_team_alias_stale_bypass_disabled_by_default for determinism - Extract _get_team_deployments helper to centralize DB access pattern - Add clarifying comments for team_public_model_name assignment ordering Made-with: Cursor * fix(router): address remaining Greptile review comments - Cache LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS at module level to avoid hot-path secret lookups - Add clarifying comments for should_include_deployment team isolation logic - Add negative assertion for update_team.assert_not_called() in test - Add docstring clarification for _get_team_deployments helper pattern - Add explicit assertion message in test_get_model_list_alias_optimization Made-with: Cursor * fix(router): address final Greptile P1/P2 comments - Reorder team_public_model_name assignment to happen before model_name mutation for clarity - Add comment explaining no-rename fast-exit case in _update_existing_team_model_assignment - Add comment explaining final patch_data.model_name = None applies to all code paths Made-with: Cursor * fix(tests): reset module-level cache in stale alias bypass tests Reset _ENABLE_TEAM_STALE_ALIAS_BYPASS to None in both test functions to ensure test isolation and prevent ordering-dependent failures Made-with: Cursor * feat(router): add order-based fallback so higher order deployments are tried on failure When order=1 deployments fail, the router now automatically tries order=2, then order=3, etc. before falling through to external fallbacks. This removes the need for enable_pre_call_checks and makes order work as a true priority-based fallback chain within a model group. Co-Authored-By: Claude Opus 4.6 * fix(router): address Greptile P0/P1 review comments on order fallback - P0: Skip order-based fallback for ContextWindowExceededError and ContentPolicyViolationError so their dedicated fallback handlers run - P1: Read _target_order from kwargs to skip already-tried order levels, preventing wasteful retries and exponential retry storms Co-Authored-By: Claude Opus 4.6 * feat(router): add order-based fallback so higher order deployments are tried on failure When a request to an order=1 deployment fails, the router now automatically tries order=2, order=3, etc. before falling through to external fallbacks. Works for all error types (429, 404, connection errors). Requires enable_pre_call_checks=True. Co-Authored-By: Claude Opus 4.6 * fix(router): handle non-standard fallback formats with order-based fallback When fallbacks use non-standard formats (e.g. ["claude-3-haiku"] or [{"model": "...", "messages": [...]}]), detect them with _check_non_standard_fallback_format and pass them through directly instead of trying to parse with get_fallback_model_group which only handles the standard dict-keyed format. Co-Authored-By: Claude Opus 4.6 * docs: remove enable_pre_call_checks requirement from order docs Order-based routing and fallback work without enable_pre_call_checks in the current code. Remove the stale requirement from both doc files. Co-Authored-By: Claude Opus 4.6 * Fix tests --------- Co-authored-by: Sameer Kankute Co-authored-by: Claude Opus 4.6 Co-authored-by: yuneng-jiang * Potential fix for code scanning alert no. 4373: Log Injection Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: Sameer Kankute Co-authored-by: Claude Opus 4.6 Co-authored-by: yuneng-jiang Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fcb8b6db80a..3ed96c163af 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -26,6 +26,22 @@ from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_head _SPECIAL_HEADERS_CACHE = frozenset( v.value.lower() for v in SpecialHeaders._member_map_.values() ) + + +def _sanitize_for_log(value: Any) -> str: + """ + Basic log sanitization helper to reduce log-injection risk. + + Removes newline and carriage-return characters so user-controlled + values cannot forge additional log lines when written to text logs. + """ + try: + text = str(value) + except Exception: + # Fallback to repr if str() fails for any reason + text = repr(value) + # Strip CR/LF characters commonly used for log injection + return text.replace("\r", "").replace("\n", "") from litellm.router import Router from litellm.secret_managers.main import get_secret_bool from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS @@ -1355,10 +1371,8 @@ def _update_model_if_team_alias_exists( "New sibling deployments may be unreachable. " "Set LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS=true to enable " "team-scoped sibling routing.", - str(_model).replace("\n", "").replace("\r", ""), - str(user_api_key_dict.team_id) - .replace("\n", "") - .replace("\r", ""), + _sanitize_for_log(_model), + user_api_key_dict.team_id, ) data["model"] = aliased_target From ab9c875a0019037d7adf59a3e14bb8625faa5eea Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 15:35:25 -0700 Subject: [PATCH 41/55] feat: scope guardrail submissions to team members Backend: - list_guardrail_submissions no longer 403s non-admins; it returns only submissions whose team_id matches one of the caller's teams (via get_user_object.teams). Admins still see all. - Filtering by a team the caller is not in returns 403. - Users with no team memberships get an empty list (no DB query). - get_guardrail_submission applies the same scoping to single-item GETs. Frontend: - Remove admin-only bail-out in TeamGuardrailsTab.fetchSubmissions so internal users actually load their team's submissions. - Finish antd migration in guardrails.tsx: drop the last Tremor Button. - Remove guardrailsList.length === 0 gate on the Test Playground tab; the playground already renders a "No guardrails available" inline empty state, which is more discoverable than a disabled tab. Tests: - Cover non-admin scoped access, empty teams, cross-team filter 403, and per-submission GET scoping. --- .../proxy/guardrails/guardrail_endpoints.py | 74 +++++++-- .../guardrails/test_guardrail_endpoints.py | 140 +++++++++++++++++- .../src/components/guardrails.tsx | 5 +- .../guardrails/TeamGuardrailsTab.tsx | 8 +- 4 files changed, 203 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index b88c6524bb7..4aa552a631d 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -724,6 +724,30 @@ def _parse_json_field(value: Any) -> Optional[Dict[str, Any]]: return None +async def _get_user_team_ids(user_api_key_dict: UserAPIKeyAuth) -> List[str]: + """Return the list of team_ids the caller belongs to (empty list if none).""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if not user_api_key_dict.user_id or prisma_client is None: + return [] + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_obj is None or not user_obj.teams: + return [] + return [t for t in user_obj.teams if t] + + def _row_to_submission_item(row: Any) -> GuardrailSubmissionItem: guardrail_info = _parse_json_field(row.guardrail_info) or {} team_guardrail = row.team_id is not None @@ -756,27 +780,49 @@ async def list_guardrail_submissions( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - List team guardrail submissions (admin only). Returns only guardrails with a team_id. + List team guardrail submissions. Returns only guardrails with a team_id. + + Admins see all submissions. Non-admin users see submissions for teams they are + a member of. Status values: pending_review (team-registered, awaiting approval), active (approved), rejected. Optional filters: - status: pending_review | active | rejected - - team_id: filter by specific team + - team_id: filter by specific team (non-admins must be a member of that team) - search: name/description """ from litellm.proxy.proxy_server import prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException(status_code=403, detail="Admin access required") - if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") + is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + visible_team_ids: Optional[List[str]] = None + if not is_admin: + visible_team_ids = await _get_user_team_ids(user_api_key_dict) + if team_id is not None and team_id not in visible_team_ids: + raise HTTPException( + status_code=403, + detail=f"You are not a member of team {team_id!r}", + ) + try: - # Single query: fetch all team guardrails (team_id is not null) + where_clause: Dict[str, Any] = {"team_id": {"not": None}} + if visible_team_ids is not None: + if not visible_team_ids: + # Non-admin with no team memberships: nothing visible. + return ListGuardrailSubmissionsResponse( + submissions=[], + summary=GuardrailSubmissionSummary( + total=0, pending_review=0, active=0, rejected=0 + ), + ) + where_clause["team_id"] = {"in": visible_team_ids} + + # Single query: fetch team guardrails visible to the caller all_team_rows = await prisma_client.db.litellm_guardrailstable.find_many( - where={"team_id": {"not": None}}, + where=where_clause, order={"created_at": "desc"}, ) @@ -837,15 +883,14 @@ async def get_guardrail_submission( guardrail_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - """Get a single guardrail submission by id (admin only).""" + """Get a single guardrail submission by id. Non-admins may only access submissions for teams they belong to.""" from litellm.proxy.proxy_server import prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException(status_code=403, detail="Admin access required") - if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") + is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + try: row = await prisma_client.db.litellm_guardrailstable.find_unique( where={"guardrail_id": guardrail_id} @@ -854,6 +899,13 @@ async def get_guardrail_submission( raise HTTPException( status_code=404, detail="Guardrail submission not found" ) + if not is_admin: + visible_team_ids = await _get_user_team_ids(user_api_key_dict) + if row.team_id is None or row.team_id not in visible_team_ids: + raise HTTPException( + status_code=403, + detail="You are not a member of the team that owns this submission", + ) return _row_to_submission_item(row) except HTTPException: raise diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index ca224726361..244150f8554 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1237,13 +1237,82 @@ async def test_register_guardrail_duplicate_name(mocker): @pytest.mark.asyncio -async def test_list_guardrail_submissions_requires_admin(mocker): - """List submissions returns 403 when user is not admin.""" +async def test_list_guardrail_submissions_non_admin_scoped_to_own_teams(mocker): + """Non-admin callers see only submissions for teams they belong to.""" + mock_prisma = mocker.Mock() + own_team_row = mocker.Mock( + guardrail_id="mine", + guardrail_name="mine-guard", + status="pending_review", + team_id="team-mine", + litellm_params={}, + guardrail_info={}, + submitted_at=None, + reviewed_at=None, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + find_many = AsyncMock(return_value=[own_team_row]) + mock_prisma.db.litellm_guardrailstable.find_many = find_many + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids", + AsyncMock(return_value=["team-mine"]), + ) + user = UserAPIKeyAuth( + user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) + + result = await list_guardrail_submissions(user_api_key_dict=user) + + # DB query scoped to visible teams + where_clause = find_many.call_args.kwargs["where"] + assert where_clause["team_id"] == {"in": ["team-mine"]} + assert len(result.submissions) == 1 + assert result.submissions[0].team_id == "team-mine" + # Summary counts reflect only visible teams + assert result.summary.total == 1 + assert result.summary.pending_review == 1 + + +@pytest.mark.asyncio +async def test_list_guardrail_submissions_non_admin_no_teams(mocker): + """Non-admin caller with no team memberships gets an empty list (not 403).""" + mock_prisma = mocker.Mock() + find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_guardrailstable.find_many = find_many + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids", + AsyncMock(return_value=[]), + ) + user = UserAPIKeyAuth( + user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) + + result = await list_guardrail_submissions(user_api_key_dict=user) + + assert result.submissions == [] + assert result.summary.total == 0 + assert find_many.call_count == 0 # no DB query when user has no teams + + +@pytest.mark.asyncio +async def test_list_guardrail_submissions_non_admin_team_filter_forbidden(mocker): + """Non-admin caller filtering by a team they're not in gets 403.""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) - user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids", + AsyncMock(return_value=["team-mine"]), + ) + user = UserAPIKeyAuth( + user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) with pytest.raises(HTTPException) as exc_info: - await list_guardrail_submissions(user_api_key_dict=user) + await list_guardrail_submissions( + team_id="team-other", user_api_key_dict=user + ) assert exc_info.value.status_code == 403 @@ -1354,6 +1423,69 @@ async def test_get_guardrail_submission_not_found(mocker): assert exc_info.value.status_code == 404 +@pytest.mark.asyncio +async def test_get_guardrail_submission_non_admin_own_team(mocker): + """Non-admin caller can fetch a submission belonging to one of their teams.""" + mock_prisma = mocker.Mock() + row = mocker.Mock( + guardrail_id="sub-1", + guardrail_name="team-guard", + status="pending_review", + team_id="team-mine", + litellm_params={}, + guardrail_info={}, + submitted_at=None, + reviewed_at=None, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids", + AsyncMock(return_value=["team-mine"]), + ) + user = UserAPIKeyAuth( + user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) + + result = await get_guardrail_submission("sub-1", user) + + assert result.guardrail_id == "sub-1" + assert result.team_id == "team-mine" + + +@pytest.mark.asyncio +async def test_get_guardrail_submission_non_admin_other_team_forbidden(mocker): + """Non-admin caller gets 403 when fetching a submission for a team they're not in.""" + mock_prisma = mocker.Mock() + row = mocker.Mock( + guardrail_id="sub-1", + guardrail_name="team-guard", + status="pending_review", + team_id="team-other", + litellm_params={}, + guardrail_info={}, + submitted_at=None, + reviewed_at=None, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids", + AsyncMock(return_value=["team-mine"]), + ) + user = UserAPIKeyAuth( + user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER + ) + + with pytest.raises(HTTPException) as exc_info: + await get_guardrail_submission("sub-1", user) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio async def test_approve_guardrail_submission_success(mocker): """Approve sets status to active and initializes guardrail in memory.""" diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index d52aba15ab0..a4383e1648f 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -1,6 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Button } from "@tremor/react"; -import { Dropdown, Tabs } from "antd"; +import { Button, Dropdown, Tabs } from "antd"; import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons"; import { getGuardrailsList, deleteGuardrailCall } from "./networking"; import AddGuardrailForm from "./guardrails/add_guardrail_form"; @@ -240,7 +239,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { key: "playground", label: "Test Playground", - disabled: !accessToken || guardrailsList.length === 0, + disabled: !accessToken, children: ( ([]); const [summary, setSummary] = useState({ total: 0, @@ -837,7 +833,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { }, [search]); const fetchSubmissions = useCallback(async () => { - if (!accessToken || !isAdmin) { + if (!accessToken) { setIsLoading(false); return; } @@ -862,7 +858,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { } finally { setIsLoading(false); } - }, [accessToken, isAdmin, statusFilter, searchDebounced]); + }, [accessToken, statusFilter, searchDebounced]); useEffect(() => { fetchSubmissions(); From 51876292a0edcefca1419794cd4bcc4661514b98 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Sat, 4 Apr 2026 16:09:42 -0700 Subject: [PATCH 42/55] Litellm ishaan april4 2 (#25150) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(router): integrate allowed_fails_policy into health check failures (#24988) * feat(router): integrate allowed_fails_policy into health check failures Health check failures now increment the same per-deployment failure counters used by allowed_fails_policy, so users can control how many health check failures of each error type are required before a deployment enters cooldown. - ahealth_check() preserves the original exception in its return dict - run_with_timeout() returns a litellm.Timeout on health check timeout - _perform_health_check() propagates exceptions to unhealthy endpoints - _write_health_state_to_router_cache() calls _set_cooldown_deployments for each unhealthy endpoint that has an exception - When allowed_fails_policy is set, the binary health check filter is bypassed so cooldown is the sole routing exclusion mechanism - Safety net: if all deployments are in cooldown with enable_health_check_routing=True, the cooldown filter is bypassed Co-Authored-By: Claude Opus 4.6 (1M context) * feat(router): add health_check_ignore_transient_errors flag When enabled, health check failures with 429 (rate limit) or 408 (timeout) status codes are skipped from the cooldown pipeline. These are transient load issues, not broken deployments. Auth errors (401), 404, and 5xx errors still increment counters and trigger cooldown as before. Config (general_settings): health_check_ignore_transient_errors: true Co-Authored-By: Claude Opus 4.6 (1M context) * fix(router): also exclude 429/408 from health state cache when ignore_transient_errors set The previous fix only skipped cooldown counter increments. The health state cache was still marking 429/408 endpoints as is_healthy=False, causing the binary health check filter to exclude them from routing. Now, when health_check_ignore_transient_errors=True, 429/408 endpoints are also excluded from the unhealthy list passed to build_deployment_health_states(), so the binary filter treats them as unaffected (not unhealthy). Co-Authored-By: Claude Opus 4.6 (1M context) * docs(router): add health check driven routing guide New standalone page covering the full health check routing feature: allowed_fails_policy integration, health_check_ignore_transient_errors, architecture SVG, step-by-step setup, and gotchas (TTL, AllowedFails semantics). Replaces the inline section in health.md with a link to the new page. Added to the Routing & Load Balancing sidebar. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(health-check-routing): fix three CI failures - Add "exception" to ILLEGAL_DISPLAY_PARAMS in health_check.py so the exception object is stripped before the health endpoint serializes results to JSON (fixes TypeError: 'URL' object is not iterable) - Add allowed_fails_policy = None to FakeRouter stubs in test_router_health_check_routing.py (fixes AttributeError) - Add health_check_ignore_transient_errors to config_settings.md router settings reference table (fixes documentation test) Co-Authored-By: Claude Opus 4.6 (1M context) * Fix litellm/tests/proxy_unit_tests/test_proxy_server.py * fix(router): address greptile review comments - Narrow cooldown safety-net bypass: only fires when allowed_fails_policy is set (cooldown is health-check driven). Without a policy, cooldowns are from real request failures and must not be bypassed. - Restore cooldown deployments DEBUG log that was accidentally removed. - Fix test_health TypeError: move exception extraction to a separate exceptions_by_model_id dict returned alongside endpoints, so exception objects never appear in the endpoint dicts that get JSON-serialized by the /health response. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(health-check-routing): properly isolate exceptions from health response Return exceptions_by_model_id as a separate third value from _perform_health_check / perform_health_check so exception objects (which contain non-JSON-serializable httpx URL types) never appear in the endpoint dicts that get serialized by the /health response. Callers updated: _health_endpoints.py, shared_health_check_manager.py, proxy_server.py background loop. All use the exceptions dict only for cooldown integration, not for display. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(shared-health-check): fix remaining 2-value return sites and update type annotation * fix(health-check-routing): fix P0 cooldown integration never firing The cooldown loop was reading endpoint.get("exception") which is always None because exceptions are now returned via exceptions_by_model_id, not stored in endpoint dicts. Fixed to use _exceptions.get(model_id). Also fixes the transient-error filter to use _exceptions instead of endpoint.get("exception"), and fixes all remaining 2-value return sites in shared_health_check_manager.py. Tests updated to pass exceptions via exceptions_by_model_id parameter instead of endpoint dicts. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(health-check-routing): fix P1 transient-error filter broken on cache hits When SharedHealthCheckManager returns cached results, exceptions_by_model_id is always {} so the transient-error filter defaulted to status 500 for all endpoints, incorrectly marking 429/408 endpoints as unhealthy. Fix: store integer exception_status on each unhealthy endpoint dict in _perform_health_check. _get_endpoint_exception_status() uses the live exception object when available (direct path) and falls back to the stored integer (cache-hit path). The integer is JSON-serializable and survives the shared cache round-trip. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(health-check-routing): gate cooldown loop behind allowed_fails_policy Without the policy, cooldown is not the routing exclusion mechanism. Firing _set_cooldown_deployments for all enable_health_check_routing users was a backwards-incompatible change — 401s would immediately cooldown deployments that the binary filter would have recovered on the next cycle. Co-Authored-By: Claude Opus 4.6 (1M context) * revert: undo allowed_fails_policy gate on cooldown loop Cooldown integration via health checks is intentional for all enable_health_check_routing users, not just those with allowed_fails_policy. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(docs+tests): fix health_check_ignore_transient_errors doc section and test coverage - Move health_check_ignore_transient_errors from router_settings to general_settings in config_settings.md (code reads it from general_settings) - Remove duplicate enable_health_check_routing / health_check_staleness_threshold entries that were incorrectly listed under router_settings - Replace TestHealthCheckEndpointExceptionPropagation tests with ones that exercise the real _perform_health_check code path via mocked ahealth_check, verifying exceptions appear in exceptions_by_model_id and NOT in endpoint dicts Co-Authored-By: Claude Opus 4.6 (1M context) * fix(tests+docs): fix tuple unpacking and docs test failures - Update test mocks that return (healthy, unhealthy) to return (healthy, unhealthy, {}) to match the new 3-value signature - Update test unpackings of perform_shared_health_check to use healthy, unhealthy, _ = ... - Add health_check_ignore_transient_errors to router_settings section in config_settings.md (it is a Router constructor param, so the doc test requires it there; it also lives in general_settings for proxy use) Co-Authored-By: Claude Opus 4.6 (1M context) * Fix CodeQL errors * fix(tests): fix 2-value unpackings of _perform_health_check in test_health_check.py * fix(tests): fix mock _perform_health_check returning 2-tuple instead of 3 * fix team routing --------- Co-authored-by: Claude Opus 4.6 (1M context) * fix: add distributed lock for key rotation job (#23364) * fix: add distributed lock for key rotation job * fix: address Greptile review feedback on key rotation lock (#23834) * fix: address Greptile review feedback on key rotation lock * fix req changes greptile * feat(proxy): Optional on_error for guardrail pipeline (API / technical failures) (#24831) * guardrails fallback * docs * docs: add LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS to environment variables reference * fix(mypy): accept Union[Dict, Any] in _get_deployment_order and use typed list to fix min() type error * fix(mypy): use Optional[str] for api_base in PydanticAI provider to match superclass signature --------- Co-authored-by: Sameer Kankute Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Co-authored-by: Shivam Rawat Co-authored-by: yuneng-jiang --- docs/my-website/docs/proxy/config_settings.md | 3 + .../proxy/guardrails/guardrail_policies.md | 4 +- .../proxy/guardrails/policy_flow_builder.md | 65 +- docs/my-website/docs/proxy/health.md | 81 +- .../docs/proxy/health_check_routing.md | 340 ++++++++ docs/my-website/sidebars.js | 3 +- litellm/constants.py | 4 + litellm/main.py | 52 +- .../common_utils/key_rotation_manager.py | 43 +- .../db_transaction_queue/pod_lock_manager.py | 15 +- litellm/proxy/health_check.py | 34 +- .../shared_health_check_manager.py | 12 +- .../health_endpoints/_health_endpoints.py | 2 +- .../proxy/policy_engine/pipeline_executor.py | 19 +- litellm/proxy/proxy_server.py | 139 ++- litellm/router.py | 150 ++-- .../proxy/policy_engine/pipeline_types.py | 16 +- litellm/utils.py | 32 +- .../litellm_utils_tests/test_health_check.py | 12 +- tests/proxy_unit_tests/test_proxy_server.py | 4 +- .../common_utils/test_key_rotation_e2e.py | 559 +++++++++++++ .../common_utils/test_key_rotation_lock.py | 229 +++++ .../policy_engine/test_pipeline_executor.py | 141 ++++ .../proxy/test_health_check_functions.py | 2 +- .../proxy/test_shared_health_check.py | 12 +- ..._health_check_allowed_fails_integration.py | 789 ++++++++++++++++++ .../test_router_health_check_routing.py | 2 + tests/test_litellm/test_router.py | 292 ++++++- .../policy_engine/test_pipeline_types.py | 14 +- .../policies/pipeline_flow_builder.tsx | 55 +- .../src/components/policies/types.ts | 2 + 31 files changed, 2873 insertions(+), 254 deletions(-) create mode 100644 docs/my-website/docs/proxy/health_check_routing.md create mode 100644 tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py create mode 100644 tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py create mode 100644 tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 27d9aed52b4..528a5c10903 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -289,6 +289,7 @@ router_settings: | database_connection_pool_timeout | integer | Database connection pool timeout in seconds | | disable_error_logs | boolean | If true, suppresses error tracking and storage in the database | | enable_health_check_routing | boolean | If true, enables health check-driven request routing to avoid unhealthy deployments | +| health_check_ignore_transient_errors | boolean | If true, 429 (rate limit) and 408 (timeout) health check failures are ignored and do not affect routing or cooldown | | enable_mcp_registry | boolean | If true, enables access to the centralized MCP server registry | | enforce_rbac | boolean | If true, enables role-based access control (RBAC) for all proxy operations | | forward_llm_provider_auth_headers | boolean | If true, forwards provider-specific auth headers to LLM API calls | @@ -397,6 +398,7 @@ router_settings: | guardrail_list | List[GuardrailTypedDict] | List of guardrail configurations for guardrail load balancing. Enables load balancing across multiple guardrail deployments with the same guardrail_name. [Further Docs](./guardrails/guardrail_load_balancing.md) | | enable_health_check_routing | boolean | If true, enables health check-driven deployment filtering to avoid routing requests to unhealthy deployments | | health_check_staleness_threshold | integer | Maximum age in seconds for cached health check results before marking deployments as stale | +| health_check_ignore_transient_errors | boolean | If true, 429 (rate limit) and 408 (timeout) health check failures are ignored and do not affect routing or cooldown | ### environment variables - Reference @@ -821,6 +823,7 @@ router_settings: | LITELLM_KEY_ROTATION_ENABLED | Enable auto-key rotation for LiteLLM (boolean). Default is false. | LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS | Interval in seconds for how often to run job that auto-rotates keys. Default is 86400 (24 hours). | LITELLM_KEY_ROTATION_GRACE_PERIOD | Duration to keep old key valid after rotation (e.g. "24h", "2d"). Default is empty (immediate revoke). Used for scheduled rotations and as fallback when not specified in regenerate request. +| LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS | TTL in seconds for the distributed lock used by the key rotation job. Default is 600 (10 minutes). | LITELLM_LICENSE | License key for LiteLLM usage | LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS | Set to `True` to use the local bundled Anthropic beta headers config only, disabling remote fetching. Default is `False` | LITELLM_LOCAL_BLOG_POSTS | When set to `True`, uses the local bundled blog posts only, disabling remote fetching from GitHub. Default is `False` diff --git a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md index f4411553c69..18c9025da6c 100644 --- a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md +++ b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md @@ -311,7 +311,7 @@ Response: ## Policy Flow Builder -For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step pass/fail actions. +For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step **pass**, **fail**, and optional **error** actions (`on_pass`, `on_fail`, `on_error`). ## Config Reference @@ -337,7 +337,7 @@ policies: | `guardrails.add` | `list[string]` | Guardrails to enable. | | `guardrails.remove` | `list[string]` | Guardrails to disable (useful with inheritance). | | `condition.model` | `string` or `list[string]` | Optional. Only apply when model matches. Supports regex. | -| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions. See [Policy Flow Builder](./policy_flow_builder). | +| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions (`on_pass`, `on_fail`, optional `on_error`). See [Policy Flow Builder](./policy_flow_builder). | ### `policy_attachments` diff --git a/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md b/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md index 2a83f3768ab..630930aa893 100644 --- a/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md +++ b/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md @@ -1,8 +1,8 @@ # Policy Flow Builder -The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail passes or fails. +The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail **passes**, **fails a policy check** (content intervention), or hits a **technical error** (e.g. timeout, unreachable provider, missing guardrail). -Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors). +Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors). With **`on_error`**, you can treat **technical** failures differently from **policy** failures—for example, fall back to another provider when the primary API errors, while still blocking on flagged content. ## When to use the Flow Builder @@ -19,6 +19,7 @@ Use the Flow Builder when you need: - **Custom responses** — return a specific message when a guardrail fails instead of a generic block - **Data chaining** — pass modified data (e.g., PII-masked content) from one step to the next - **Fine-grained control** — different actions on pass vs. fail per step +- **Technical-error routing** — set `on_error` separately from `on_fail` so outages or timeouts can **allow**, **block**, **go to the next step**, or return a **custom response** without conflating them with content violations ## Concepts @@ -29,24 +30,37 @@ A pipeline has: - **Mode**: `pre_call` (before the LLM) or `post_call` (after the LLM) - **Steps**: Ordered list of guardrail steps +### Outcomes: pass, fail, and error + +Each step run produces one of three outcomes: + +| Outcome | Meaning | Typical cause | +|--------|---------|----------------| +| **pass** | Guardrail completed without blocking | Content allowed, or data was modified and returned | +| **fail** | Policy intervention | Guardrail raised an intervention (e.g. flagged content, blocked request) | +| **error** | Technical failure | Timeouts, network errors, guardrail not registered, or other non-intervention exceptions | + +`on_pass` and `on_fail` apply to **pass** and **fail** respectively. **`on_error`** applies only to **error**. If `on_error` is omitted, the pipeline uses **`on_fail`** for error outcomes (backward compatible). + ### Step actions -Each step defines what happens when the guardrail **passes** and when it **fails**: +For each step you choose an action for **pass**, **fail**, and optionally **error**. Allowed values are: `next`, `allow`, `block`, `modify_response`. | Action | Description | |--------|-------------| -| **Next Step** | Continue to the next guardrail in the pipeline | -| **Allow** | Stop the pipeline and allow the request to proceed | -| **Block** | Stop the pipeline and block the request | -| **Custom Response** | Return a custom message instead of the default block | +| **Next Step** (`next`) | Continue to the next guardrail in the pipeline | +| **Allow** (`allow`) | Stop the pipeline and allow the request to proceed | +| **Block** (`block`) | Stop the pipeline and block the request | +| **Custom Response** (`modify_response`) | Return a custom message instead of the default block | ### Step options | Field | Type | Description | |-------|------|--------------| | `guardrail` | `string` | Name of the guardrail to run | -| `on_pass` | `string` | Action when guardrail passes: `next`, `allow`, `block`, `modify_response` | -| `on_fail` | `string` | Action when guardrail fails: `next`, `allow`, `block`, `modify_response` | +| `on_pass` | `string` | Action when outcome is **pass**: `next`, `allow`, `block`, `modify_response` | +| `on_fail` | `string` | Action when outcome is **fail** (policy intervention): `next`, `allow`, `block`, `modify_response` | +| `on_error` | `string` (optional) | Action when outcome is **error** (technical). If omitted, **error** uses `on_fail`. | | `pass_data` | `boolean` | Forward modified request data (e.g., PII-masked) to the next step | | `modify_response_message` | `string` | Custom message when using `modify_response` action | @@ -57,7 +71,7 @@ Each step defines what happens when the guardrail **passes** and when it **fails 3. Select **Flow Builder** (instead of the simple form) 4. Design your flow: - **Trigger** — Incoming LLM request (runs when the policy matches) - - **Steps** — Add guardrails, set ON PASS and ON FAIL actions per step + - **Steps** — Add guardrails, set **ON PASS**, **ON FAIL**, and **ON ERROR** actions per step (ON ERROR is optional; when unset, errors follow ON FAIL) - **End** — Request proceeds to the LLM 5. Use the **+** between steps to insert new steps 6. Use the **Test** panel to run sample messages through the pipeline before saving @@ -151,6 +165,37 @@ policies: First attempt passes → allow. First attempt fails → retry the same guardrail; second pass → allow, second fail → block. +## Technical errors vs policy failures (`on_error`) + +Use **`on_error`** when you want different behavior for **API/infra problems** than for **content policy** violations. + +- **`on_fail`** — Runs when the guardrail **intervenes** (e.g. toxic content, PII detected). +- **`on_error`** — Runs when the step ends in **error** (timeout, connection failure, guardrail not loaded, etc.). If you omit `on_error`, **error** outcomes use **`on_fail`**. + +Example: block on bad content, but if the primary scanner is down, fall back to a second guardrail instead of blocking every request: + +```yaml +policies: + error-fallback-policy: + guardrails: + add: + - primary_scanner + - backup_scanner + pipeline: + mode: pre_call + steps: + - guardrail: primary_scanner + on_pass: allow + on_fail: block + on_error: next + - guardrail: backup_scanner + on_pass: allow + on_fail: block + on_error: allow +``` + +If `primary_scanner` errors → run `backup_scanner`. If `backup_scanner` errors → allow the request (set `on_error` to `block` if you prefer fail-closed). + ## Example: Custom response on fail Return a branded message instead of a generic block: diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 530bea3d06b..1d893961b62 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -316,86 +316,9 @@ general_settings: ## Health Check Driven Routing -By default, background health checks are observability-only — they populate the `/health` endpoint but don't affect routing. Unhealthy deployments still receive traffic until request failures trigger cooldown. +Route traffic away from unhealthy deployments proactively — before user requests hit them. Supports per-error-type failure thresholds, transient error suppression, and automatic safety nets. -With `enable_health_check_routing: true`, the router **excludes deployments that failed their last background health check** before selecting a candidate. This gives you proactive failover instead of reactive cooldown. - -### How it works - -1. Background health checks run on their configured interval -2. After each cycle, every deployment is marked healthy or unhealthy -3. On each incoming request, the router filters out unhealthy deployments **before** cooldown filtering and load balancing -4. If all deployments are unhealthy, the filter is bypassed (safety net — never causes a total outage) -5. If health state is stale (older than `health_check_staleness_threshold`), it is ignored - -### Quick start - -```yaml -model_list: - - model_name: gpt-4 - litellm_params: - model: openai/gpt-4 - api_key: os.environ/OPENAI_API_KEY - - model_name: gpt-4 - litellm_params: - model: openai/gpt-4 - api_key: os.environ/OPENAI_API_KEY_SECONDARY - -general_settings: - background_health_checks: true - health_check_interval: 60 - enable_health_check_routing: true -``` - -### Configuration - -| Setting | Where | Default | Description | -|---------|-------|---------|-------------| -| `enable_health_check_routing` | `general_settings` | `false` | Enable/disable health-check-driven routing | -| `health_check_staleness_threshold` | `general_settings` | `health_check_interval * 2` | Seconds before health state is considered stale and ignored | -| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work | -| `health_check_interval` | `general_settings` | `300` | Seconds between health check cycles | - -### Interaction with cooldown - -Health check filtering and cooldown are **additive**. A deployment can be excluded by either mechanism: - -- **Health check filter** — proactive, runs on the configured interval, excludes deployments that failed the last check -- **Cooldown** — reactive, triggered by request failures, excludes deployments for a short TTL - -This means request failures still provide fast detection between health check intervals. - -### Staleness - -If a health check result is older than `health_check_staleness_threshold`, it is ignored and the deployment is treated as eligible. This prevents stale data from permanently excluding a deployment if the health check loop stops or slows down. - -The default staleness threshold is `health_check_interval * 2`. For a 60s interval, health state expires after 120s. - -### Example: custom staleness - -```yaml -general_settings: - background_health_checks: true - health_check_interval: 30 - enable_health_check_routing: true - health_check_staleness_threshold: 90 # ignore health state older than 90s -``` - -### Debugging - -Run the proxy with `--detailed_debug` and look for: - -``` -health_check_routing_state_updated healthy=3 unhealthy=1 -``` - -This is logged after each health check cycle when routing state is written. - -If the safety net triggers (all deployments unhealthy), you'll see: - -``` -All deployments marked unhealthy by health checks, bypassing health filter -``` +See the full guide: [Health Check Driven Routing](./health_check_routing.md) ## Health Check Timeout diff --git a/docs/my-website/docs/proxy/health_check_routing.md b/docs/my-website/docs/proxy/health_check_routing.md new file mode 100644 index 00000000000..daf0b19212c --- /dev/null +++ b/docs/my-website/docs/proxy/health_check_routing.md @@ -0,0 +1,340 @@ +# Health Check Driven Routing + +Route traffic away from unhealthy deployments before users hit errors. Background health checks run on a configurable interval, and any deployment that fails gets removed from the routing pool proactively, not after a user request already failed. + + +## Architecture + + + {/* Background */} + + + {/* LEFT PANEL: Background health check loop */} + + Background Loop + every health_check_interval seconds + + {/* Deployment A */} + + Deployment A + ahealth_check() → 200 ✓ + + {/* Deployment B */} + + Deployment B + ahealth_check() → 401 ✗ + + {/* Deployment C */} + + Deployment C + ahealth_check() → 429 ⚡ + + {/* ignore_transient box */} + + ignore_transient_errors: true + 429 / 408 → ignored + not written to cache + + {/* allowed_fails_policy box */} + + allowed_fails_policy + 401 → increment counter + counter > threshold + → cooldown triggered + + {/* CENTER PANEL: Shared State */} + + Shared State + + {/* Health State Cache */} + + DeploymentHealthCache + A → healthy ✓ + B → unhealthy ✗ + C → not written (ignored) + TTL: staleness_threshold × 1.5 + + {/* Cooldown Cache */} + + Cooldown Cache + B → cooling down + (after policy threshold) + TTL: cooldown_time + + {/* failed_calls counter */} + + failed_calls counter + B: 2 / AuthAllowedFails: 1 + → threshold exceeded + TTL: cooldown_time (must > interval) + + {/* RIGHT PANEL: Request path */} + + Request Path + + {/* Incoming request */} + + Incoming request + + {/* All deployments */} + + All deployments [A, B, C] + + + + {/* Health check filter */} + + ① Health Check Filter + if policy set → bypass + else → remove unhealthy + + + + {/* Cooldown filter */} + + ② Cooldown Filter + remove deployments in cooldown + + + + {/* Safety net */} + + Safety Net + if all removed → return all + + + + {/* Load balancer */} + + ③ Load Balancer + + + + {/* Selected deployment */} + + Selected: Deployment A ✓ + + + + {/* ARROWS: left → center */} + + + + + + {/* ARROWS: center → right */} + + + + {/* Arrow markers */} + + + + + + + + + + + + + + + + + + + + + + + +## What problem does this solve? + +By default, LiteLLM routes traffic to all deployments and only stops sending to a broken one after it has already failed a user request. The cooldown system is reactive. + +Health check driven routing makes this **proactive**: a background loop pings every deployment on a configurable interval. If a deployment fails its health check, it gets removed from the routing pool immediately, before a user request lands on it. + +When you also set `allowed_fails_policy`, you control exactly how many health check failures of each error type (auth errors, rate limits, timeouts) are needed before a deployment enters cooldown. This avoids false positives from transient noise. + + +## Setup + +### Step 1: Enable background health checks + +Background health checks are off by default. Turn them on in `general_settings`: + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 60 # seconds between each full check cycle +``` + +### Step 2: Enable health check routing + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 60 + enable_health_check_routing: true # ← route away from unhealthy deployments +``` + +At this point, any deployment that fails its health check is immediately excluded from routing until the next check cycle clears it. + +### Step 3: Add a policy to control how many failures trigger cooldown + +Without a policy, the first health check failure marks a deployment as unhealthy. If you want more tolerance (e.g., only act after 2 consecutive auth failures), use `allowed_fails_policy`: + +```yaml +model_list: + - model_name: claude-sonnet + litellm_params: + model: anthropic/claude-sonnet-4-5 + api_key: os.environ/ANTHROPIC_API_KEY + + - model_name: claude-sonnet + litellm_params: + model: anthropic/claude-sonnet-4-5 + api_key: os.environ/ANTHROPIC_API_KEY_SECONDARY + +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + +router_settings: + cooldown_time: 60 # how long a deployment stays in cooldown + allowed_fails_policy: + AuthenticationErrorAllowedFails: 1 # cooldown after 2nd auth failure + TimeoutErrorAllowedFails: 3 # cooldown after 4th timeout +``` + +When `allowed_fails_policy` is set, the binary health check filter is bypassed. Only the cooldown system controls routing exclusion, and it only fires after your configured threshold is crossed. + +### Step 4 (optional): Ignore transient errors + +429 (rate limit) and 408 (timeout) from a health check usually mean the deployment is temporarily overloaded, not broken. To prevent these from affecting routing at all: + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + health_check_ignore_transient_errors: true # 429 and 408 never affect routing +``` + +With this on, only hard failures (401, 404, 5xx) from health checks contribute to cooldown. + + +## Full example + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY_SECONDARY + + - model_name: gpt-4o + litellm_params: + model: azure/gpt-4o + api_base: os.environ/AZURE_API_BASE + api_key: os.environ/AZURE_API_KEY + +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + health_check_ignore_transient_errors: true + +router_settings: + cooldown_time: 60 + allowed_fails_policy: + AuthenticationErrorAllowedFails: 0 # cooldown immediately on auth failure + TimeoutErrorAllowedFails: 2 # cooldown after 3 timeouts + RateLimitErrorAllowedFails: 5 # cooldown after 6 rate limits (if not ignoring transients) +``` + + +## Configuration reference + +| Setting | Where | Default | Description | +|---|---|---|---| +| `enable_health_check_routing` | `general_settings` | `false` | Route away from deployments that fail health checks | +| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work | +| `health_check_interval` | `general_settings` | `300` | Seconds between full health check cycles | +| `health_check_staleness_threshold` | `general_settings` | `interval x 2` | Seconds before cached health state is ignored | +| `health_check_ignore_transient_errors` | `general_settings` | `false` | Ignore 429 and 408 from health checks; these never affect routing | +| `cooldown_time` | `router_settings` | `5` | Seconds a deployment stays in cooldown after threshold is crossed | +| `allowed_fails_policy` | `router_settings` | `null` | Per-error-type failure thresholds before cooldown (see below) | + +### `allowed_fails_policy` fields + +| Field | Error type | HTTP status | +|---|---|---| +| `AuthenticationErrorAllowedFails` | Bad API key | 401 | +| `TimeoutErrorAllowedFails` | Request timeout | 408 | +| `RateLimitErrorAllowedFails` | Rate limit exceeded | 429 | +| `BadRequestErrorAllowedFails` | Malformed request | 400 | +| `ContentPolicyViolationErrorAllowedFails` | Content filtered | 400 | + +The value is the number of failures **tolerated** before cooldown. `0` means cooldown on the first failure. `2` means cooldown on the third. + + +## Things to keep in mind + +- **Counter TTL must be longer than the health check interval.** `allowed_fails_policy` works by incrementing a `failed_calls` counter per deployment. That counter expires after `cooldown_time` seconds. If `cooldown_time` is shorter than `health_check_interval`, the counter resets between every check cycle and failures never accumulate. Set `cooldown_time` greater than `health_check_interval` when using `allowed_fails_policy`. + + ```yaml + router_settings: + cooldown_time: 60 # must be > health_check_interval (30s here) + + general_settings: + health_check_interval: 30 + ``` + +- **`AllowedFails: N` means cooldown on the (N+1)th failure.** The counter check is `updated_fails > allowed_fails`, so `0` triggers on the 1st failure, `1` on the 2nd, `2` on the 3rd. + + | `AllowedFails` | Cooldown triggers after | + |---|---| + | `0` | 1st failure | + | `1` | 2nd failure | + | `2` | 3rd failure | + +- **Without `allowed_fails_policy`, the first failure is enough.** The first failed health check immediately excludes the deployment from routing. Use `allowed_fails_policy` when you want tolerance for flaky checks. + +- **If all deployments are unhealthy, the filter is bypassed.** Traffic keeps flowing rather than returning no deployment at all. Requests will fail, but the router keeps trying. + +- **Health check failures and request failures share the same counters.** When `allowed_fails_policy` is set, both sources increment the same `failed_calls` counter. A deployment at 1 health check failure that then receives 1 failing request will hit the threshold for `AllowedFails: 1` and enter cooldown. + + +## Debugging + +Run the proxy with `--detailed_debug` and look for these log lines: + +After each health check cycle (written at DEBUG level): +``` +health_check_routing_state_updated healthy=2 unhealthy=1 +``` + +When a health check failure increments the counter and triggers cooldown (DEBUG level): +``` +checks 'should_run_cooldown_logic' +Attempting to add to cooldown list +``` + +When safety net fires because all deployments are in cooldown: +``` +All deployments in cooldown via health-check routing, bypassing cooldown filter +``` + +When safety net fires because all deployments are unhealthy (binary filter, no `allowed_fails_policy`): +``` +All deployments marked unhealthy by health checks, bypassing health filter +``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index b514ea2234c..ab4ae46ec3b 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -1051,7 +1051,8 @@ const sidebars = { "proxy/fallback_management", "proxy/tag_routing", "proxy/timeout", - "wildcard_routing" + "wildcard_routing", + "proxy/health_check_routing" ], }, { diff --git a/litellm/constants.py b/litellm/constants.py index 252068bd7b0..a9facabb010 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1319,6 +1319,9 @@ LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int( LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv( "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" ) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default) +LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( + os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) +) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" @@ -1347,6 +1350,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000) ) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" +KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/main.py b/litellm/main.py index cbedd1735c7..ddd37b47536 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3792,9 +3792,9 @@ def completion( # type: ignore # noqa: PLR0915 "aws_region_name" not in optional_params or optional_params["aws_region_name"] is None ): - optional_params[ - "aws_region_name" - ] = aws_bedrock_client.meta.region_name + optional_params["aws_region_name"] = ( + aws_bedrock_client.meta.region_name + ) bedrock_route = BedrockModelInfo.get_bedrock_route(model) if bedrock_route == "converse": @@ -6214,9 +6214,9 @@ def adapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[ - Union[BaseModel, AdapterCompletionStreamWrapper] - ] = None + translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = ( + None + ) if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params( response=response @@ -6396,9 +6396,9 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params[ - "audio_transcription_duration" - ] = calculated_duration + response._hidden_params["audio_transcription_duration"] = ( + calculated_duration + ) return response except Exception as e: @@ -6621,9 +6621,9 @@ def transcription( if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params[ - "audio_transcription_duration" - ] = calculated_duration + response._hidden_params["audio_transcription_duration"] = ( + calculated_duration + ) if response is None: raise ValueError("Unmapped provider passed in. Unable to get the response.") @@ -6927,9 +6927,9 @@ def speech( # noqa: PLR0915 ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY ] = query_params - litellm_params_dict[ - ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY - ] = voice_id + litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = ( + voice_id + ) if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -7250,7 +7250,8 @@ async def ahealth_check( if mode is None: return { - "error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}" + "error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}", + "exception": e, } error_to_return = str(e) + "\nstack trace: " + stack_trace @@ -7262,6 +7263,7 @@ async def ahealth_check( return { "error": error_to_return, "raw_request_typed_dict": raw_request_typed_dict, + "exception": e, } @@ -7508,9 +7510,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(content_chunks) > 0: - response["choices"][0]["message"][ - "content" - ] = processor.get_combined_content(content_chunks) + response["choices"][0]["message"]["content"] = ( + processor.get_combined_content(content_chunks) + ) thinking_blocks = [ chunk @@ -7521,9 +7523,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(thinking_blocks) > 0: - response["choices"][0]["message"][ - "thinking_blocks" - ] = processor.get_combined_thinking_content(thinking_blocks) + response["choices"][0]["message"]["thinking_blocks"] = ( + processor.get_combined_thinking_content(thinking_blocks) + ) reasoning_chunks = [ chunk @@ -7534,9 +7536,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(reasoning_chunks) > 0: - response["choices"][0]["message"][ - "reasoning_content" - ] = processor.get_combined_reasoning_content(reasoning_chunks) + response["choices"][0]["message"]["reasoning_content"] = ( + processor.get_combined_reasoning_content(reasoning_chunks) + ) annotation_chunks = [ chunk diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 5a0a1fabc7d..aaf39a7a19d 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, LITELLM_KEY_ROTATION_GRACE_PERIOD, + LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, ) from litellm.proxy._types import ( GenerateKeyResponse, @@ -30,14 +31,42 @@ class KeyRotationManager: Manages automated key rotation based on individual key rotation schedules. """ - def __init__(self, prisma_client: PrismaClient): + def __init__(self, prisma_client: PrismaClient, pod_lock_manager=None): self.prisma_client = prisma_client + self.pod_lock_manager = pod_lock_manager async def process_rotations(self): """ - Main entry point - find and rotate keys that are due for rotation + Main entry point - find and rotate keys that are due for rotation. + Uses PodLockManager to ensure only one pod runs rotation in multi-pod deployments. """ + from litellm.constants import KEY_ROTATION_JOB_NAME + + lock_acquired = False try: + # If we have a pod lock manager with Redis, try to acquire the lock + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + # Use a dedicated lock TTL (default 600s) instead of the check interval + # (which defaults to 86400s / 24h). Using the check interval would create + # a 24-hour deadlock window if a pod crashes before releasing the lock. + lock_ttl = max( + LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, 300 + ) # At least 5 minutes, configurable via LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS + lock_acquired = ( + await self.pod_lock_manager.acquire_lock( + cronjob_id=KEY_ROTATION_JOB_NAME, + ttl=lock_ttl, + ) + or False + ) + if not lock_acquired: + verbose_proxy_logger.warning( + "Key rotation: another pod is already running rotation " + "or Redis lock acquisition failed — skipping this cycle. " + "Keys will be rotated on the next cycle." + ) + return + verbose_proxy_logger.info("Starting scheduled key rotation check...") # Clean up expired deprecated keys first @@ -74,6 +103,16 @@ class KeyRotationManager: except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") + finally: + # Only release the lock if it was actually acquired + if ( + lock_acquired + and self.pod_lock_manager + and self.pod_lock_manager.redis_cache + ): + await self.pod_lock_manager.release_lock( + cronjob_id=KEY_ROTATION_JOB_NAME, + ) async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 546ea05998c..6435498ae03 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -32,6 +32,7 @@ class PodLockManager: async def acquire_lock( self, cronjob_id: str, + ttl: Optional[int] = None, ) -> Optional[bool]: """ Attempt to acquire the lock for a specific cron job using Redis. @@ -39,15 +40,20 @@ class PodLockManager: Args: cronjob_id: The ID of the cron job to lock + ttl: Optional custom TTL in seconds. Defaults to DEFAULT_CRON_JOB_LOCK_TTL_SECONDS. + Use a longer TTL for jobs that may take longer than the default 60s + (e.g. key rotation with many keys). """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") return None try: + lock_ttl = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS verbose_proxy_logger.debug( - "Pod %s attempting to acquire Redis lock for cronjob_id=%s", + "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", self.pod_id, cronjob_id, + lock_ttl, ) # Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX) # and with an expiration (EX) to avoid deadlocks. @@ -56,7 +62,7 @@ class PodLockManager: lock_key, self.pod_id, nx=True, - ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, + ttl=lock_ttl, ) if acquired: verbose_proxy_logger.info( @@ -133,11 +139,10 @@ class PodLockManager: ) else: verbose_proxy_logger.warning( - "Spend tracking - pod %s failed to release Redis lock for cronjob_id=%s. " - "Lock will expire after TTL=%ds.", + "Pod %s failed to release Redis lock for cronjob_id=%s. " + "Lock will expire after its TTL.", self.pod_id, cronjob_id, - DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, ) else: verbose_proxy_logger.debug( diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 3e05ee3c484..5d1bcf31f84 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -21,6 +21,8 @@ ILLEGAL_DISPLAY_PARAMS = [ "vertex_credentials", "aws_access_key_id", "aws_secret_access_key", + "exception", # internal; not JSON-serializable, never for display + "litellm_metadata", # internal tracking metadata with auth objects; not for display ] MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"] @@ -95,7 +97,12 @@ async def run_with_timeout(task, timeout): except asyncio.TimeoutError: # `asyncio.wait_for()` already cancels only the awaited task on timeout. # Do not cancel unrelated sibling health check tasks. - return {"error": "Timeout exceeded"} + timeout_exception = litellm.Timeout( + message="Health check timeout exceeded", + model="", + llm_provider="", + ) + return {"error": "Timeout exceeded", "exception": timeout_exception} async def _run_model_health_check(model: dict): @@ -204,6 +211,10 @@ async def _perform_health_check( healthy_endpoints = [] unhealthy_endpoints = [] + # Exceptions keyed by model_id; returned separately so callers can use + # them for cooldown integration without risking JSON-serialization errors + # in the /health response. + exceptions_by_model_id: dict = {} for is_healthy, model in zip(results, model_list): litellm_params = model["litellm_params"] @@ -218,14 +229,23 @@ async def _perform_health_check( cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details) if _model_id: cleaned["model_id"] = _model_id + if "exception" in is_healthy: + exc = is_healthy["exception"] + exceptions_by_model_id[_model_id] = exc + # Store integer status code so shared-cache readers can + # reconstruct the transient-error filter without the exception object. + cleaned["exception_status"] = getattr(exc, "status_code", 500) unhealthy_endpoints.append(cleaned) else: cleaned = _clean_endpoint_data(litellm_params, details) if _model_id: cleaned["model_id"] = _model_id + if isinstance(is_healthy, Exception): + exceptions_by_model_id[_model_id] = is_healthy + cleaned["exception_status"] = getattr(is_healthy, "status_code", 500) unhealthy_endpoints.append(cleaned) - return healthy_endpoints, unhealthy_endpoints + return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id def build_deployment_health_states( @@ -366,7 +386,7 @@ async def perform_health_check( source, cycle_id, ) - return [], [] + return [], [], {} cycle_start_time = time.monotonic() requested_model_count = len(model_list) @@ -406,7 +426,11 @@ async def perform_health_check( ) try: - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + ( + healthy_endpoints, + unhealthy_endpoints, + exceptions_by_model_id, + ) = await _perform_health_check( model_list, details, max_concurrency=max_concurrency, @@ -438,4 +462,4 @@ async def perform_health_check( _rss_mb_for_log(), ) - return healthy_endpoints, unhealthy_endpoints + return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index ae18a42c02b..2ecee5095b8 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -192,7 +192,7 @@ class SharedHealthCheckManager: model_list: List[Dict[str, Any]], details: bool = True, max_concurrency: Optional[int] = None, - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, Any]]: """ Perform health check with shared state coordination. @@ -217,6 +217,7 @@ class SharedHealthCheckManager: return ( cached_results.get("healthy_endpoints", []), cached_results.get("unhealthy_endpoints", []), + {}, ) # No recent cache, try to acquire lock @@ -231,7 +232,11 @@ class SharedHealthCheckManager: len(model_list), ) - healthy_endpoints, unhealthy_endpoints = await perform_health_check( + ( + healthy_endpoints, + unhealthy_endpoints, + exceptions_by_model_id, + ) = await perform_health_check( model_list=model_list, details=details, max_concurrency=max_concurrency, @@ -242,7 +247,7 @@ class SharedHealthCheckManager: healthy_endpoints, unhealthy_endpoints ) - return healthy_endpoints, unhealthy_endpoints + return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id finally: # Always release the lock @@ -262,6 +267,7 @@ class SharedHealthCheckManager: return ( cached_results.get("healthy_endpoints", []), cached_results.get("unhealthy_endpoints", []), + {}, ) # Still no cache, fall back to local health check diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index ef9436f2d8c..8a09edfd4c4 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -771,7 +771,7 @@ async def _perform_health_check_and_save( max_concurrency=None, ): """Helper function to perform health check and save results to database""" - healthy_endpoints, unhealthy_endpoints = await perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check( model_list=model_list, cli_model=cli_model, model=target_model, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 729b42ce638..3c5a1d67be4 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -74,7 +74,7 @@ class PipelineExecutor: duration = time.perf_counter() - start_time - action = step.on_pass if outcome == "pass" else step.on_fail + action = _pipeline_action_for_outcome(step, outcome) step_result = PipelineStepResult( guardrail_name=step.guardrail, @@ -206,6 +206,23 @@ class PipelineExecutor: return None +def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: + """ + Map pipeline step outcome to the configured action. + + - pass -> on_pass + - fail -> on_fail (content/policy intervention) + - error -> on_error if set, else on_fail (backward compatible) + """ + if outcome == "pass": + return step.on_pass + if outcome == "fail": + return step.on_fail + if step.on_error is not None: + return step.on_error + return step.on_fail + + def _extract_error_message(e: Exception) -> str: """Extract a human-readable error message from a guardrail exception.""" if isinstance(e, ModifyResponseException): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2a661a8348e..fd88f44fbcf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -635,9 +635,9 @@ except ImportError: server_root_path = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1532,9 +1532,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1545,13 +1545,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[ - str -] = [] # Models that should use native provider background mode instead of polling +native_background_mode: List[str] = ( + [] +) # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -2036,9 +2036,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -2248,23 +2248,54 @@ def _schedule_background_health_check_db_save( ) +def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int: + """Return the HTTP status code for an unhealthy endpoint. + + Prefers the live exception object in `exceptions` (direct health check path). + Falls back to the `exception_status` integer stored on the endpoint dict + (shared-cache path, where exception objects are not available). + """ + model_id = endpoint.get("model_id") + exc = exceptions.get(model_id) if model_id else None + if exc is not None: + return getattr(exc, "status_code", 500) + return endpoint.get("exception_status", 500) + + def _write_health_state_to_router_cache( healthy_endpoints: list, unhealthy_endpoints: list, + exceptions_by_model_id: Optional[dict] = None, ) -> None: """ Write deployment health states to the router's health state cache for health-check-driven routing. No-op if the feature is disabled. """ from litellm.proxy.health_check import build_deployment_health_states + from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + increment_deployment_failures_for_current_minute, + ) + + _exceptions: dict = exceptions_by_model_id or {} try: if llm_router is None or not llm_router.enable_health_check_routing: return + # When health_check_ignore_transient_errors is set, treat 429/408 + # endpoints as healthy so they are not filtered from routing. + _effective_unhealthy = unhealthy_endpoints + if llm_router.health_check_ignore_transient_errors: + _effective_unhealthy = [ + ep + for ep in unhealthy_endpoints + if _get_endpoint_exception_status(ep, _exceptions) not in (429, 408) + ] + states = build_deployment_health_states( healthy_endpoints=healthy_endpoints, - unhealthy_endpoints=unhealthy_endpoints, + unhealthy_endpoints=_effective_unhealthy, ) if states: llm_router.health_state_cache.set_deployment_health_states(states) @@ -2273,6 +2304,37 @@ def _write_health_state_to_router_cache( sum(1 for s in states.values() if s.get("is_healthy")), sum(1 for s in states.values() if not s.get("is_healthy")), ) + + for endpoint in unhealthy_endpoints: + model_id = endpoint.get("model_id") + if not model_id: + continue + + original_exception = _exceptions.get(model_id) + if original_exception is None: + continue + + exception_status = getattr(original_exception, "status_code", 500) + + if llm_router.health_check_ignore_transient_errors and exception_status in ( + 429, + 408, + ): + continue + + increment_deployment_failures_for_current_minute( + litellm_router_instance=llm_router, + deployment_id=model_id, + ) + + _set_cooldown_deployments( + litellm_router_instance=llm_router, + original_exception=original_exception, + exception_status=exception_status, + deployment=model_id, + time_to_cooldown=llm_router.cooldown_time, + ) + except Exception as e: verbose_proxy_logger.warning( "Failed to write health state to router cache: %s", str(e) @@ -2384,6 +2446,7 @@ async def _run_background_health_check(): ( healthy_endpoints, unhealthy_endpoints, + _exceptions_by_model_id, ) = await shared_health_manager.perform_shared_health_check( model_list=_llm_model_list, details=details_bool, @@ -2397,6 +2460,7 @@ async def _run_background_health_check(): ( healthy_endpoints, unhealthy_endpoints, + _exceptions_by_model_id, ) = await _run_direct_health_check_with_instrumentation( _llm_model_list, health_check_details, @@ -2407,6 +2471,7 @@ async def _run_background_health_check(): ( healthy_endpoints, unhealthy_endpoints, + _exceptions_by_model_id, ) = await _run_direct_health_check_with_instrumentation( _llm_model_list, health_check_details, @@ -2449,7 +2514,9 @@ async def _run_background_health_check(): ) # Write health state to router cache for health-check-driven routing - _write_health_state_to_router_cache(healthy_endpoints, unhealthy_endpoints) + _write_health_state_to_router_cache( + healthy_endpoints, unhealthy_endpoints, _exceptions_by_model_id + ) await asyncio.sleep(health_check_interval) @@ -3245,6 +3312,7 @@ class ProxyConfig: general_settings = {} _enable_hc_routing = False _hc_staleness = None + _hc_ignore_transient = False if general_settings: ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### key_management_settings = general_settings.get( @@ -3437,6 +3505,9 @@ class ProxyConfig: _hc_staleness = general_settings.get( "health_check_staleness_threshold", None ) + _hc_ignore_transient = general_settings.get( + "health_check_ignore_transient_errors", False + ) verbose_proxy_logger.info( "background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s health_check_routing=%s", use_background_health_checks, @@ -3479,6 +3550,8 @@ class ProxyConfig: router_params["enable_health_check_routing"] = True if _hc_staleness is not None: router_params["health_check_staleness_threshold"] = _hc_staleness + if _hc_ignore_transient: + router_params["health_check_ignore_transient_errors"] = True ## MODEL LIST model_list = config.get("model_list", None) if model_list: @@ -5217,10 +5290,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -5602,9 +5675,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -6503,10 +6576,18 @@ class ProxyStartupEvent: KeyRotationManager, ) - # Get prisma_client from global scope + # Get prisma_client and proxy_logging_obj from global scope global prisma_client + global proxy_logging_obj if prisma_client is not None: - key_rotation_manager = KeyRotationManager(prisma_client) + # Reuse the PodLockManager from db_spend_update_writer + pod_lock_manager = ( + proxy_logging_obj.db_spend_update_writer.pod_lock_manager + ) + key_rotation_manager = KeyRotationManager( + prisma_client, + pod_lock_manager=pod_lock_manager, + ) verbose_proxy_logger.debug( f"Key rotation background job scheduled every {LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS} seconds (LITELLM_KEY_ROTATION_ENABLED=true)" ) @@ -12624,9 +12705,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/litellm/router.py b/litellm/router.py index 1b8f7c91761..a58b3ce25e1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -310,6 +310,7 @@ class Router: ignore_invalid_deployments: bool = False, enable_health_check_routing: bool = False, health_check_staleness_threshold: Optional[int] = None, + health_check_ignore_transient_errors: bool = False, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -408,9 +409,9 @@ class Router: ) # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### - cache_type: Literal[ - "local", "redis", "redis-semantic", "s3", "disk" - ] = "local" # default to an in-memory cache + cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = ( + "local" # default to an in-memory cache + ) redis_cache = None cache_config: Dict[str, Any] = {} @@ -458,9 +459,9 @@ class Router: self.default_max_parallel_requests = default_max_parallel_requests self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() - self.team_pattern_routers: Dict[ - str, PatternMatchRouter - ] = {} # {"TEAM_ID": PatternMatchRouter} + self.team_pattern_routers: Dict[str, PatternMatchRouter] = ( + {} + ) # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} self.complexity_routers: Dict[str, "ComplexityRouter"] = {} @@ -501,6 +502,7 @@ class Router: ) self.disable_cooldowns = disable_cooldowns self.enable_health_check_routing = enable_health_check_routing + self.health_check_ignore_transient_errors = health_check_ignore_transient_errors _staleness = health_check_staleness_threshold or ( DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER ) @@ -655,12 +657,12 @@ class Router: ) ) - self.model_group_retry_policy: Optional[ - Dict[str, RetryPolicy] - ] = model_group_retry_policy - self.model_group_affinity_config: Optional[ - Dict[str, List[str]] - ] = model_group_affinity_config + self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = ( + model_group_retry_policy + ) + self.model_group_affinity_config: Optional[Dict[str, List[str]]] = ( + model_group_affinity_config + ) self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None if allowed_fails_policy is not None: @@ -2066,7 +2068,10 @@ class Router: async def _acompletion( # noqa: PLR0915 self, model: str, messages: List[Dict[str, str]], **kwargs - ) -> Union[ModelResponse, CustomStreamWrapper,]: + ) -> Union[ + ModelResponse, + CustomStreamWrapper, + ]: """ - Get an available deployment - call it with a semaphore over the call @@ -4300,9 +4305,9 @@ class Router: healthy_deployments=healthy_deployments, responses=responses ) returned_response = cast(OpenAIFileObject, responses[0]) - returned_response._hidden_params[ - "model_file_id_mapping" - ] = model_file_id_mapping + returned_response._hidden_params["model_file_id_mapping"] = ( + model_file_id_mapping + ) return returned_response except Exception as e: verbose_router_logger.exception( @@ -5310,11 +5315,16 @@ class Router: e, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError), ) - all_deployments = self._get_all_deployments(model_name=original_model_group) + _request_team_id: Optional[str] = ( + kwargs.get("metadata", {}) or {} + ).get("user_api_key_team_id") + all_deployments = self._get_all_deployments( + model_name=original_model_group, team_id=_request_team_id + ) _order_set: set = { - d.get("litellm_params", {}).get("order") + litellm.utils._get_deployment_order(d) for d in all_deployments - if d.get("litellm_params", {}).get("order") is not None + if litellm.utils._get_deployment_order(d) is not None } order_values: list = sorted(_order_set) if len(order_values) > 1 and not _skip_order_fallback: @@ -5387,11 +5397,11 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: - context_window_fallback_model_group: Optional[ - List[str] - ] = self._get_fallback_model_group_from_fallbacks( - fallbacks=context_window_fallbacks, - model_group=model_group, + context_window_fallback_model_group: Optional[List[str]] = ( + self._get_fallback_model_group_from_fallbacks( + fallbacks=context_window_fallbacks, + model_group=model_group, + ) ) if context_window_fallback_model_group is None: raise original_exception @@ -5423,11 +5433,11 @@ class Router: e.message += "\n{}".format(error_message) elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: - content_policy_fallback_model_group: Optional[ - List[str] - ] = self._get_fallback_model_group_from_fallbacks( - fallbacks=content_policy_fallbacks, - model_group=model_group, + content_policy_fallback_model_group: Optional[List[str]] = ( + self._get_fallback_model_group_from_fallbacks( + fallbacks=content_policy_fallbacks, + model_group=model_group, + ) ) if content_policy_fallback_model_group is None: raise original_exception @@ -5649,9 +5659,9 @@ class Router: ) ## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking _metadata["attempted_retries"] = 0 - _metadata[ - "max_retries" - ] = num_retries # Updated after overrides in exception handler + _metadata["max_retries"] = ( + num_retries # Updated after overrides in exception handler + ) try: self._handle_mock_testing_rate_limit_error( model_group=model_group, kwargs=kwargs @@ -6770,26 +6780,26 @@ class Router: """ from litellm.router_strategy.auto_router.auto_router import AutoRouter - auto_router_config_path: Optional[ - str - ] = deployment.litellm_params.auto_router_config_path + auto_router_config_path: Optional[str] = ( + deployment.litellm_params.auto_router_config_path + ) auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config if auto_router_config_path is None and auto_router_config is None: raise ValueError( "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" ) - default_model: Optional[ - str - ] = deployment.litellm_params.auto_router_default_model + default_model: Optional[str] = ( + deployment.litellm_params.auto_router_default_model + ) if default_model is None: raise ValueError( "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" ) - embedding_model: Optional[ - str - ] = deployment.litellm_params.auto_router_embedding_model + embedding_model: Optional[str] = ( + deployment.litellm_params.auto_router_embedding_model + ) if embedding_model is None: raise ValueError( "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" @@ -6832,13 +6842,13 @@ class Router: ComplexityRouter, ) - complexity_router_config: Optional[ - dict - ] = deployment.litellm_params.complexity_router_config + complexity_router_config: Optional[dict] = ( + deployment.litellm_params.complexity_router_config + ) - default_model: Optional[ - str - ] = deployment.litellm_params.complexity_router_default_model + default_model: Optional[str] = ( + deployment.litellm_params.complexity_router_default_model + ) # If no default model specified, try to get from config tiers if default_model is None and complexity_router_config: @@ -7497,9 +7507,9 @@ class Router: # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: - credentials[ - "custom_llm_provider" - ] = deployment.litellm_params.custom_llm_provider + credentials["custom_llm_provider"] = ( + deployment.litellm_params.custom_llm_provider + ) elif "/" in deployment.litellm_params.model: # Extract provider from "provider/model" format credentials["custom_llm_provider"] = deployment.litellm_params.model.split( @@ -9070,7 +9080,9 @@ class Router: ## get healthy deployments ### get all deployments - healthy_deployments = self._get_all_deployments(model_name=model) + healthy_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) if len(healthy_deployments) == 0: # check if the user sent in a deployment name instead @@ -9091,7 +9103,9 @@ class Router: ) # Re-assign model to the fallback and try to get deployments again model = fallback_model - healthy_deployments = self._get_all_deployments(model_name=model) + healthy_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) # If still no deployments after checking for fallbacks, raise an error if len(healthy_deployments) == 0: @@ -9181,10 +9195,23 @@ class Router: ) if verbose_router_logger.isEnabledFor(logging.DEBUG): verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") + _pre_cooldown_deployments = healthy_deployments healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, ) + # Safety net: only bypass cooldown filter when health-check routing is + # driving cooldown (i.e. allowed_fails_policy is set). Without a policy, + # cooldowns are from real request failures and must not be bypassed. + if ( + not healthy_deployments + and self.enable_health_check_routing + and self.allowed_fails_policy is not None + ): + verbose_router_logger.warning( + "All deployments in cooldown via health-check routing, bypassing cooldown filter" + ) + healthy_deployments = _pre_cooldown_deployments healthy_deployments = await self.async_callback_filter_deployments( model=model, @@ -9617,10 +9644,20 @@ class Router: cooldown_deployments = _get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) + _pre_cooldown_deployments = healthy_deployments healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, ) + if ( + not healthy_deployments + and self.enable_health_check_routing + and self.allowed_fails_policy is not None + ): + verbose_router_logger.warning( + "All deployments in cooldown via health-check routing, bypassing cooldown filter" + ) + healthy_deployments = _pre_cooldown_deployments # filter pre-call checks if self.enable_pre_call_checks and messages is not None: @@ -9922,6 +9959,12 @@ class Router: if not self.enable_health_check_routing: return healthy_deployments + # When allowed_fails_policy is set, cooldown is the sole routing exclusion + # mechanism -- skip the binary health check filter so the policy threshold + # is respected before any deployment is excluded. + if self.allowed_fails_policy is not None: + return healthy_deployments + unhealthy_ids = ( await self.health_state_cache.async_get_unhealthy_deployment_ids( parent_otel_span=parent_otel_span @@ -9951,6 +9994,9 @@ class Router: if not self.enable_health_check_routing: return healthy_deployments + if self.allowed_fails_policy is not None: + return healthy_deployments + unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids( parent_otel_span=parent_otel_span ) diff --git a/litellm/types/proxy/policy_engine/pipeline_types.py b/litellm/types/proxy/policy_engine/pipeline_types.py index 29d2e576000..abbb127cd7a 100644 --- a/litellm/types/proxy/policy_engine/pipeline_types.py +++ b/litellm/types/proxy/policy_engine/pipeline_types.py @@ -18,18 +18,24 @@ class PipelineStep(BaseModel): """ A single step in a guardrail pipeline. - Each step runs a guardrail and takes an action based on pass/fail. + Each step runs a guardrail and takes an action based on pass, policy fail, + or technical/API error (see pipeline executor outcome types). """ guardrail: str = Field(description="Name of the guardrail to run.") on_fail: str = Field( default="block", - description="Action when guardrail rejects: next | block | allow | modify_response", + description="Action when guardrail rejects content (policy intervention): next | block | allow | modify_response", ) on_pass: str = Field( default="allow", description="Action when guardrail passes: next | block | allow | modify_response", ) + on_error: Optional[str] = Field( + default=None, + description="Action when the guardrail raises a technical error (timeouts, " + "unreachable provider, non-intervention HTTP errors). If omitted, uses on_fail.", + ) pass_data: bool = Field( default=False, description="Forward modified request data (e.g., PII-masked) to next step.", @@ -41,9 +47,11 @@ class PipelineStep(BaseModel): model_config = ConfigDict(extra="forbid") - @field_validator("on_fail", "on_pass") + @field_validator("on_fail", "on_pass", "on_error") @classmethod - def validate_action(cls, v: str) -> str: + def validate_action(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return None if v not in VALID_PIPELINE_ACTIONS: raise ValueError( f"Invalid action '{v}'. Must be one of: {sorted(VALID_PIPELINE_ACTIONS)}" diff --git a/litellm/utils.py b/litellm/utils.py index 6806961bf51..f902644e760 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4875,6 +4875,19 @@ def calculate_max_parallel_requests( return None +def _get_deployment_order(deployment: Union[Dict, Any]) -> Optional[int]: + """ + Returns the routing order for a deployment. + + Checks litellm_params first (static config), then model_info (dynamic/team + models added via API where order lives in model_info, not litellm_params). + """ + order = deployment.get("litellm_params", {}).get("order") + if order is None: + order = deployment.get("model_info", {}).get("order") + return order + + def _get_order_filtered_deployments( healthy_deployments: List[Dict], target_order: Optional[int] = None ) -> List: @@ -4882,7 +4895,7 @@ def _get_order_filtered_deployments( filtered = [ d for d in healthy_deployments - if d["litellm_params"].get("order") == target_order + if _get_deployment_order(d) == target_order ] if filtered: return filtered @@ -4890,20 +4903,19 @@ def _get_order_filtered_deployments( return healthy_deployments # Default: pick min order group - min_order = min( - ( - deployment["litellm_params"]["order"] - for deployment in healthy_deployments - if "order" in deployment["litellm_params"] - ), - default=None, - ) + _valid_orders: List[int] = [ + o + for deployment in healthy_deployments + for o in [_get_deployment_order(deployment)] + if o is not None + ] + min_order: Optional[int] = min(_valid_orders) if _valid_orders else None if min_order is not None: filtered_deployments = [ deployment for deployment in healthy_deployments - if deployment["litellm_params"].get("order") == min_order + if _get_deployment_order(deployment) == min_order ] return filtered_deployments diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 0b0e091211f..708b2403c49 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -475,13 +475,13 @@ async def test_perform_health_check_filters_by_model_id(): captured_list.append(m_list) return [ {"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]} - ], [] + ], [], {} with patch( "litellm.proxy.health_check._perform_health_check", side_effect=mock_perform_health_check, ): - healthy_endpoints, unhealthy_endpoints = await perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check( model_list=model_list, model_id="deployment-id-2", details=True ) @@ -521,7 +521,7 @@ async def test_perform_health_check_with_health_check_model(): return {"status": "healthy"} with patch("litellm.ahealth_check", side_effect=mock_health_check): - healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list) + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list) print("health check calls: ", health_check_calls) # Verify the health check used the override model @@ -556,7 +556,7 @@ async def test_health_check_bad_model(): }, ] details = None - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check( model_list, details ) print(f"healthy_endpoints: {healthy_endpoints}") @@ -574,7 +574,7 @@ async def test_health_check_bad_model(): "litellm.ahealth_check", side_effect=mock_health_check ) as mock_health_check: start_time = time.time() - healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list) + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list) end_time = time.time() print("health check calls: ", health_check_calls) assert len(healthy_endpoints) == 0 @@ -667,7 +667,7 @@ async def test_timeout_does_not_cancel_other_health_checks(): return {"status": "healthy"} with patch("litellm.ahealth_check", side_effect=mock_health_check): - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check( model_list, max_concurrency=1 ) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 047193055d8..7da4d41fbf1 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2420,7 +2420,7 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) - return (["healthy"], ["unhealthy"]) + return (["healthy"], ["unhealthy"], {}) monkeypatch.setattr(proxy_server, "health_check_interval", 1) monkeypatch.setattr(proxy_server, "health_check_details", None) @@ -2471,7 +2471,7 @@ async def test_background_health_check_skip_disabled_models(monkeypatch): async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) - return (["healthy"], []) + return (["healthy"], [], {}) monkeypatch.setattr(proxy_server, "health_check_interval", 1) monkeypatch.setattr(proxy_server, "health_check_details", None) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py new file mode 100644 index 00000000000..f6ef02a86de --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py @@ -0,0 +1,559 @@ +""" +End-to-end tests for key rotation feature. + +Covers the critical gaps: +1. Multi-pod simulation: two KeyRotationManagers sharing one PodLockManager +2. Error resilience: partial failures, regenerate_key_fn failures, hook failures +3. Full process_rotations flow with actual key finding + rotation + lock +4. Initialization wiring: PodLockManager is correctly passed +5. Multiple keys: some succeed, some fail, all are attempted +6. Rotation count increments correctly over multiple rotations +""" + +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + GenerateKeyResponse, + LiteLLM_VerificationToken, +) +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestMultiPodKeyRotation: + """ + Simulate two pods sharing one Redis lock to verify only one pod + runs key rotation at a time. + """ + + @pytest.mark.asyncio + async def test_two_pods_only_one_rotates(self): + """ + Two KeyRotationManagers with separate pod_lock_managers but + the same Redis backend. Only the first to acquire the lock + should rotate; the second should skip. + """ + mock_prisma = AsyncMock() + + # Shared state to simulate Redis SET NX behavior + redis_lock = {"holder": None} + + async def make_acquire_lock(pod_id): + async def acquire(cronjob_id, **kwargs): + if redis_lock["holder"] is None: + redis_lock["holder"] = pod_id + return True + return redis_lock["holder"] == pod_id + + return acquire + + async def make_release_lock(pod_id): + async def release(cronjob_id): + if redis_lock["holder"] == pod_id: + redis_lock["holder"] = None + + return release + + # Pod A + pod_a_lock_mgr = MagicMock() + pod_a_lock_mgr.redis_cache = MagicMock() + pod_a_lock_mgr.acquire_lock = AsyncMock( + side_effect=await make_acquire_lock("pod-a") + ) + pod_a_lock_mgr.release_lock = AsyncMock( + side_effect=await make_release_lock("pod-a") + ) + + # Pod B + pod_b_lock_mgr = MagicMock() + pod_b_lock_mgr.redis_cache = MagicMock() + pod_b_lock_mgr.acquire_lock = AsyncMock( + side_effect=await make_acquire_lock("pod-b") + ) + pod_b_lock_mgr.release_lock = AsyncMock( + side_effect=await make_release_lock("pod-b") + ) + + manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock_mgr) + manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock_mgr) + + # Both share the same mock methods for rotation logic + for mgr in [manager_a, manager_b]: + mgr._cleanup_expired_deprecated_keys = AsyncMock() + mgr._find_keys_needing_rotation = AsyncMock(return_value=[]) + + # Pod A acquires lock first + await manager_a.process_rotations() + # Pod A should have run rotation + manager_a._cleanup_expired_deprecated_keys.assert_called_once() + manager_a._find_keys_needing_rotation.assert_called_once() + + # Lock is released after pod A finishes, so pod B can now acquire + # But let's simulate pod B trying WHILE pod A holds the lock + # Reset the lock state to simulate concurrent access + redis_lock["holder"] = "pod-a" # Pod A holds the lock + + await manager_b.process_rotations() + # Pod B should NOT have run rotation (lock held by pod-a) + manager_b._cleanup_expired_deprecated_keys.assert_not_called() + manager_b._find_keys_needing_rotation.assert_not_called() + + @pytest.mark.asyncio + async def test_second_pod_runs_after_first_releases(self): + """ + After the first pod releases the lock, the second pod should + be able to acquire and run rotation. + """ + mock_prisma = AsyncMock() + + call_order = [] + + # Pod A - always gets the lock + pod_a_lock = MagicMock() + pod_a_lock.redis_cache = MagicMock() + pod_a_lock.acquire_lock = AsyncMock(return_value=True) + pod_a_lock.release_lock = AsyncMock() + + # Pod B - also gets the lock (simulating after A releases) + pod_b_lock = MagicMock() + pod_b_lock.redis_cache = MagicMock() + pod_b_lock.acquire_lock = AsyncMock(return_value=True) + pod_b_lock.release_lock = AsyncMock() + + manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock) + manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock) + + async def cleanup_a(): + call_order.append("a_cleanup") + + async def cleanup_b(): + call_order.append("b_cleanup") + + manager_a._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_a) + manager_a._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager_b._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_b) + manager_b._find_keys_needing_rotation = AsyncMock(return_value=[]) + + # Run sequentially: A then B + await manager_a.process_rotations() + await manager_b.process_rotations() + + # Both should have run + assert call_order == ["a_cleanup", "b_cleanup"] + pod_a_lock.release_lock.assert_called_once() + pod_b_lock.release_lock.assert_called_once() + + +class TestKeyRotationErrorResilience: + """ + Tests that key rotation handles errors gracefully: + - regenerate_key_fn failure for one key doesn't block others + - Hook failure doesn't crash the process + - Database update failure is handled + """ + + @pytest.mark.asyncio + async def test_one_key_fails_others_still_rotate(self): + """ + If rotation fails for one key, the remaining keys should still + be attempted. No key should be silently skipped. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key1 = LiteLLM_VerificationToken( + token="token-1", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-1", + ) + key2 = LiteLLM_VerificationToken( + token="token-2", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-2", + ) + key3 = LiteLLM_VerificationToken( + token="token-3", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-3", + ) + + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[key1, key2, key3]) + + rotate_calls = [] + + async def mock_rotate(key): + rotate_calls.append(key.token) + if key.token == "token-2": + raise Exception("Database connection lost") + + manager._rotate_key = AsyncMock(side_effect=mock_rotate) + + await manager.process_rotations() + + # All 3 keys should have been attempted + assert rotate_calls == ["token-1", "token-2", "token-3"] + + @pytest.mark.asyncio + async def test_regenerate_key_fn_failure_is_caught(self): + """ + If regenerate_key_fn throws, _rotate_key should propagate the error + but process_rotations should catch it per-key. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="test-key", + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + side_effect=Exception("regenerate failed: DB timeout"), + ): + # _rotate_key should raise + with pytest.raises(Exception, match="regenerate failed"): + await manager._rotate_key(key) + + # But process_rotations should catch per-key errors + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[key]) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + side_effect=Exception("regenerate failed: DB timeout"), + ): + # Should NOT raise - error is caught per-key + await manager.process_rotations() + + @pytest.mark.asyncio + async def test_hook_failure_does_not_prevent_db_update(self): + """ + If the rotation hook (async_key_rotated_hook) fails, the database + update for rotation_count should still have succeeded (it runs before the hook). + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + mock_response = GenerateKeyResponse( + key="new-key", token_id="new-token-id", user_id="test-user" + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + side_effect=Exception("Hook failed: secret manager down"), + ): + # This will raise because the hook fails + with pytest.raises(Exception, match="Hook failed"): + await manager._rotate_key(key) + + # The DB update should have been called BEFORE the hook + mock_prisma.db.litellm_verificationtoken.update.assert_called_once() + update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ + "data" + ] + assert update_data["rotation_count"] == 1 + + @pytest.mark.asyncio + async def test_cleanup_failure_does_not_prevent_rotation(self): + """ + If deprecated key cleanup fails, the rotation should still proceed. + """ + mock_prisma = AsyncMock() + mock_pod_lock = MagicMock() + mock_pod_lock.redis_cache = MagicMock() + mock_pod_lock.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_pod_lock) + + # Cleanup fails + manager._cleanup_expired_deprecated_keys = AsyncMock( + side_effect=Exception("Deprecated table doesn't exist") + ) + + # process_rotations catches the exception internally (try/except), + # but the lock must still be released in the finally block. + await manager.process_rotations() + + # Lock should still be released in finally block + mock_pod_lock.release_lock.assert_called_once() + + +class TestKeyRotationFullFlow: + """ + Full end-to-end flow tests: find keys -> rotate -> update DB -> release lock + """ + + @pytest.mark.asyncio + async def test_full_rotation_flow_with_lock(self): + """ + Full flow: acquire lock -> cleanup -> find keys -> rotate -> update DB -> release lock + """ + mock_prisma = AsyncMock() + + # Setup lock manager + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + key = LiteLLM_VerificationToken( + token="old-token-hash", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=datetime.now(timezone.utc) - timedelta(seconds=60), + rotation_count=2, + key_name="my-key", + key_alias="prod/my-key", + ) + + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + user_id="system", + ) + + # Mock cleanup + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 1 + # Mock find keys + mock_prisma.db.litellm_verificationtoken.find_many.return_value = [key] + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager.process_rotations() + + # Verify full flow executed: + # 1. Lock acquired + mock_lock.acquire_lock.assert_called_once() + + # 2. Cleanup ran + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.assert_called_once() + + # 3. Keys were queried + mock_prisma.db.litellm_verificationtoken.find_many.assert_called_once() + + # 4. DB was updated with new rotation info + mock_prisma.db.litellm_verificationtoken.update.assert_called_once() + update_args = mock_prisma.db.litellm_verificationtoken.update.call_args[1] + assert update_args["where"]["token"] == "new-token-hash" + assert update_args["data"]["rotation_count"] == 3 # was 2, now 3 + + # 5. Lock released + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_rotation_count_increments_across_multiple_rotations(self): + """ + Simulate 3 consecutive rotations and verify rotation_count increments + correctly each time: 0 -> 1 -> 2 -> 3 + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + rotation_counts_seen = [] + + for expected_count in range(3): + key = LiteLLM_VerificationToken( + token=f"token-v{expected_count}", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=expected_count, + ) + + mock_response = GenerateKeyResponse( + key=f"sk-new-v{expected_count + 1}", + token_id=f"token-v{expected_count + 1}", + user_id="system", + ) + + mock_prisma.db.litellm_verificationtoken.update.reset_mock() + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager._rotate_key(key) + + update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ + "data" + ] + rotation_counts_seen.append(update_data["rotation_count"]) + + assert rotation_counts_seen == [1, 2, 3] + + @pytest.mark.asyncio + async def test_no_keys_to_rotate_skips_gracefully(self): + """ + When no keys need rotation, process should complete without errors. + """ + mock_prisma = AsyncMock() + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 0 + mock_prisma.db.litellm_verificationtoken.find_many.return_value = [] + + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + await manager.process_rotations() + + # Verify no rotation was attempted + mock_prisma.db.litellm_verificationtoken.update.assert_not_called() + # But lock was still properly released + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_regenerate_response_missing_token_id_skips_db_update(self): + """ + If regenerate_key_fn returns a response without token_id, + the DB update for rotation metadata should be skipped. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="old-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + # Response with no token_id + mock_response = GenerateKeyResponse( + key="sk-new", + token_id=None, + user_id="system", + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager._rotate_key(key) + + # DB update should NOT have been called (no token_id) + mock_prisma.db.litellm_verificationtoken.update.assert_not_called() + + +class TestKeyRotationInitialization: + """ + Tests that the PodLockManager wiring in proxy_server.py is correct. + """ + + @pytest.mark.asyncio + async def test_key_rotation_manager_receives_pod_lock_manager(self): + """ + Verify KeyRotationManager stores the pod_lock_manager correctly. + """ + mock_prisma = AsyncMock() + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + assert manager.pod_lock_manager is mock_lock + assert manager.prisma_client is mock_prisma + + @pytest.mark.asyncio + async def test_key_rotation_manager_default_no_lock(self): + """ + When no pod_lock_manager is provided, it defaults to None. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + assert manager.pod_lock_manager is None + + @pytest.mark.asyncio + async def test_lock_pattern_matches_spend_log_cleanup(self): + """ + Verify the key rotation lock pattern is identical to spend_log_cleanup: + - acquire_lock with cronjob_id + - release_lock in finally + - lock_acquired flag guards release + """ + mock_prisma = AsyncMock() + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + + await manager.process_rotations() + + # Pattern check: acquire with cronjob_id + acquire_call = mock_lock.acquire_lock.call_args + assert "cronjob_id" in acquire_call.kwargs or len(acquire_call.args) > 0 + + # Pattern check: release with same cronjob_id + release_call = mock_lock.release_lock.call_args + assert "cronjob_id" in release_call.kwargs or len(release_call.args) > 0 + + # Both should use the same job name + from litellm.constants import KEY_ROTATION_JOB_NAME + + assert acquire_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME + assert release_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py new file mode 100644 index 00000000000..c0b3611b2b4 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py @@ -0,0 +1,229 @@ +""" +Test distributed lock behavior for key rotation manager. + +Verifies that PodLockManager is correctly used to prevent concurrent +key rotation across multiple pods in a distributed deployment. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestKeyRotationLock: + """Test distributed lock behavior in KeyRotationManager.""" + + @pytest.mark.asyncio + async def test_process_rotations_acquires_lock(self): + """ + When PodLockManager is provided and lock is acquired, + rotation logic should run normally. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() # Redis is available + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Mock _find_keys_needing_rotation to return empty list (no keys to rotate) + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was acquired with custom TTL + mock_pod_lock_manager.acquire_lock.assert_called_once() + call_kwargs = mock_pod_lock_manager.acquire_lock.call_args + assert call_kwargs.kwargs["cronjob_id"] == "litellm_key_rotation_job" + assert call_kwargs.kwargs["ttl"] >= 300 # At least 5 minutes + + # Verify rotation logic ran (cleanup + find keys called) + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + # Verify lock was released + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_skips_when_lock_held(self): + """ + When lock is held by another pod, process_rotations() should + return early without performing any rotation. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock() + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was attempted + mock_pod_lock_manager.acquire_lock.assert_called_once() + + # Verify rotation logic was NOT executed + manager._cleanup_expired_deprecated_keys.assert_not_called() + manager._find_keys_needing_rotation.assert_not_called() + + # Verify lock was NOT released (since it was never acquired) + mock_pod_lock_manager.release_lock.assert_not_called() + + @pytest.mark.asyncio + async def test_process_rotations_releases_lock_on_success(self): + """ + Lock should be released in the finally block after successful rotation. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Simulate finding and rotating a key successfully + mock_key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="test-key", + ) + manager._find_keys_needing_rotation = AsyncMock(return_value=[mock_key]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._rotate_key = AsyncMock() + + await manager.process_rotations() + + # Verify rotation was performed + manager._rotate_key.assert_called_once_with(mock_key) + + # Verify lock was released after success + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_releases_lock_on_error(self): + """ + Lock should be released in the finally block even if rotation + throws an exception. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Simulate an error during cleanup + manager._cleanup_expired_deprecated_keys = AsyncMock( + side_effect=Exception("Database connection failed") + ) + + await manager.process_rotations() + + # Verify lock was still released despite the error + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_works_without_lock_manager(self): + """ + When pod_lock_manager=None, rotation should run normally + without any lock logic (backward compat / single-pod mode). + """ + mock_prisma_client = AsyncMock() + + # No pod_lock_manager provided (default None) + manager = KeyRotationManager(mock_prisma_client) + + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify rotation logic ran normally + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + @pytest.mark.asyncio + async def test_process_rotations_works_without_redis_cache(self): + """ + When pod_lock_manager exists but redis_cache is None (no Redis configured), + rotation should run normally without locking. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = None # No Redis available + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was NOT attempted (no Redis) + mock_pod_lock_manager.acquire_lock.assert_not_called() + + # Verify rotation logic still ran + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + @pytest.mark.asyncio + async def test_process_rotations_handles_none_lock_result(self): + """ + When acquire_lock returns None (edge case), it should be treated + as lock NOT acquired, and rotation should be skipped. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=None) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock() + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify rotation logic was NOT executed (None treated as False via `or False`) + manager._cleanup_expired_deprecated_keys.assert_not_called() + manager._find_keys_needing_rotation.assert_not_called() + + # Verify lock was NOT released (lock_acquired is False) + mock_pod_lock_manager.release_lock.assert_not_called() diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 226e88bea3e..ffe8947fc61 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -46,6 +46,28 @@ class AlwaysFailGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="Content policy violation") +class HttpStatusGuardrail(CustomGuardrail): + """Raises HTTPException with a configurable status (e.g. 503 for API outage).""" + + def __init__(self, guardrail_name: str, status_code: int): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + ) + self.status_code = status_code + self.calls = 0 + + def should_run_guardrail(self, data, event_type) -> bool: + return True + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + raise HTTPException( + status_code=self.status_code, detail="Simulated HTTP error" + ) + + class AlwaysPassGuardrail(CustomGuardrail): """Mock guardrail that always passes.""" @@ -350,6 +372,125 @@ async def test_guardrail_not_found_uses_on_fail(): litellm.callbacks = original_callbacks +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(): + """ + Policy intervention (400) uses on_fail; technical error (503) uses on_error. + + Primary returns 503 -> on_error: next -> fallback runs -> allow. + """ + primary = HttpStatusGuardrail("primary-mod", status_code=503) + fallback = AlwaysPassGuardrail("fallback-filter") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep( + guardrail="primary-mod", + on_fail="block", + on_error="next", + on_pass="allow", + ), + PipelineStep( + guardrail="fallback-filter", + on_fail="block", + on_pass="allow", + ), + ], + ) + + original_callbacks = litellm.callbacks.copy() + litellm.callbacks = [primary, fallback] + + try: + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "any"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="mod-fallback", + ) + + assert primary.calls == 1 + assert fallback.calls == 1 + assert result.terminal_action == "allow" + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "next" + assert result.step_results[1].outcome == "pass" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_on_fail_next_on_content_on_error_block_stops_api_fallback(): + """ + Content policy fail (400) uses on_fail: next; API error uses on_error: block (no second step). + """ + primary_content = AlwaysFailGuardrail("strict-mod") + primary_api = HttpStatusGuardrail("strict-mod", status_code=503) + fallback = AlwaysPassGuardrail("fallback-filter") + + # Content violation: on_fail next -> would reach fallback if we had two steps + pipeline_content = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep( + guardrail="strict-mod", + on_fail="next", + on_error="block", + on_pass="allow", + ), + PipelineStep( + guardrail="fallback-filter", + on_fail="block", + on_pass="allow", + ), + ], + ) + + original_callbacks = litellm.callbacks.copy() + litellm.callbacks = [primary_content, fallback] + + try: + result = await PipelineExecutor.execute_steps( + steps=pipeline_content.steps, + mode=pipeline_content.mode, + data={"messages": [{"role": "user", "content": "bad"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) + assert result.terminal_action == "allow" + assert primary_content.calls == 1 + assert fallback.calls == 1 + finally: + litellm.callbacks = original_callbacks + + # API outage: on_error block -> do not run fallback + fallback.calls = 0 + original_callbacks = litellm.callbacks.copy() + litellm.callbacks = [primary_api, fallback] + try: + result = await PipelineExecutor.execute_steps( + steps=pipeline_content.steps, + mode=pipeline_content.mode, + data={"messages": [{"role": "user", "content": "ok"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) + assert result.terminal_action == "block" + assert primary_api.calls == 1 + assert fallback.calls == 0 + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "block" + finally: + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_guardrail_not_found_with_next_continues(): """ diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 354698b02fe..13d2131efad 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -481,7 +481,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c unhealthy = [] async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None, max_concurrency=None): - return healthy, unhealthy + return healthy, unhealthy, {} with patch( "litellm.proxy.health_endpoints._health_endpoints.perform_health_check", diff --git a/tests/test_litellm/proxy/test_shared_health_check.py b/tests/test_litellm/proxy/test_shared_health_check.py index 0212d87baab..20c96c8152d 100644 --- a/tests/test_litellm/proxy/test_shared_health_check.py +++ b/tests/test_litellm/proxy/test_shared_health_check.py @@ -246,7 +246,7 @@ class TestSharedHealthCheckManager: model_list = [{"model_name": "test-model", "litellm_params": {"model": "test-model"}}] with patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform: - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) @@ -268,9 +268,9 @@ class TestSharedHealthCheckManager: expected_unhealthy = [] with patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform: - mock_perform.return_value = (expected_healthy, expected_unhealthy) + mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) @@ -302,7 +302,7 @@ class TestSharedHealthCheckManager: model_list = [{"model_name": "test-model", "litellm_params": {"model": "test-model"}}] with patch("asyncio.sleep") as mock_sleep: # Mock sleep to avoid actual delay - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) @@ -324,9 +324,9 @@ class TestSharedHealthCheckManager: with patch("asyncio.sleep") as mock_sleep, \ patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform: - mock_perform.return_value = (expected_healthy, expected_unhealthy) + mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py new file mode 100644 index 00000000000..e2c13b952dd --- /dev/null +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -0,0 +1,789 @@ +""" +Tests for health check failures integrating with allowed_fails_policy cooldown pipeline. + +When enable_health_check_routing is True and a health check fails, the failure +should increment the same counters used by allowed_fails_policy, using the +actual exception type from the health check error. +""" + +from unittest.mock import patch + +import pytest + +import litellm +from litellm.proxy.health_check import run_with_timeout +from litellm.router import Router +from litellm.types.router import AllowedFailsPolicy + + +def _make_model(model_id: str, model_name: str = "gpt-4") -> dict: + return { + "model_name": model_name, + "litellm_params": {"model": model_name, "api_key": "fake-key"}, + "model_info": {"id": model_id}, + } + + +class TestAhealthCheckExceptionPreservation: + """Test that ahealth_check() preserves the exception object in its return dict.""" + + @pytest.mark.asyncio + async def test_run_with_timeout_returns_timeout_exception(self): + """run_with_timeout should return a litellm.Timeout in the 'exception' key on timeout.""" + import asyncio + + async def slow_task(): + await asyncio.sleep(10) + + result = await run_with_timeout(slow_task(), timeout=0.01) + + assert "error" in result + assert "exception" in result + assert isinstance(result["exception"], litellm.Timeout) + + +class TestHealthCheckEndpointExceptionPropagation: + """Test that _perform_health_check returns exceptions via exceptions_by_model_id.""" + + @pytest.mark.asyncio + async def test_unhealthy_endpoint_dict_exception_in_map(self): + """When ahealth_check returns {"error": ..., "exception": e}, the exception + must appear in exceptions_by_model_id keyed by model_id — not in the endpoint dict.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy.health_check import _perform_health_check + + auth_error = litellm.AuthenticationError( + message="Invalid key", llm_provider="openai", model="gpt-4" + ) + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"id": "deploy-abc"}, + } + ] + + with patch( + "litellm.proxy.health_check.litellm.ahealth_check", + new=AsyncMock(return_value={"error": "auth failed", "exception": auth_error}), + ): + healthy, unhealthy, exc_map = await _perform_health_check(model_list) + + assert len(unhealthy) == 1 + assert "exception" not in unhealthy[0], "exception must not be in endpoint dict" + assert exc_map.get("deploy-abc") is auth_error + + @pytest.mark.asyncio + async def test_raw_exception_from_gather_in_map(self): + """When asyncio.gather returns a raw Exception, it must appear in + exceptions_by_model_id — not in the endpoint dict.""" + from unittest.mock import patch + + from litellm.proxy.health_check import _perform_health_check + + raw_exc = litellm.RateLimitError( + message="Rate limited", llm_provider="openai", model="gpt-4" + ) + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"id": "deploy-xyz"}, + } + ] + + # Simulate asyncio.gather returning a raw exception for this task + with patch( + "litellm.proxy.health_check._run_model_health_check", + side_effect=raw_exc, + ): + healthy, unhealthy, exc_map = await _perform_health_check(model_list) + + assert len(unhealthy) == 1 + assert "exception" not in unhealthy[0], "exception must not be in endpoint dict" + assert exc_map.get("deploy-xyz") is raw_exc + + +class TestGetAllowedFailsFromPolicyWithHealthCheckExceptions: + """Test that get_allowed_fails_from_policy correctly resolves thresholds for health-check exceptions.""" + + @pytest.mark.parametrize( + "exception_type, policy_field, threshold", + [ + (litellm.Timeout, "TimeoutErrorAllowedFails", 5), + (litellm.AuthenticationError, "AuthenticationErrorAllowedFails", 3), + (litellm.RateLimitError, "RateLimitErrorAllowedFails", 10), + ( + litellm.ContentPolicyViolationError, + "ContentPolicyViolationErrorAllowedFails", + 2, + ), + (litellm.BadRequestError, "BadRequestErrorAllowedFails", 7), + ], + ) + def test_policy_resolves_for_health_check_exception_types( + self, exception_type, policy_field, threshold + ): + """Each exception type from a health check should resolve to its policy threshold.""" + policy = AllowedFailsPolicy(**{policy_field: threshold}) + router = Router( + model_list=[_make_model("d1")], + allowed_fails_policy=policy, + ) + exception = exception_type( + message="health check failed", llm_provider="openai", model="gpt-4" + ) + result = router.get_allowed_fails_from_policy(exception=exception) + assert result == threshold + + def test_policy_returns_none_for_unmatched_exception(self): + """When no policy field matches the exception type, return None (fall back to allowed_fails).""" + policy = AllowedFailsPolicy(TimeoutErrorAllowedFails=5) + router = Router( + model_list=[_make_model("d1")], + allowed_fails_policy=policy, + ) + # Use a generic Exception that doesn't match any policy field + result = router.get_allowed_fails_from_policy(exception=Exception("generic")) + assert result is None + + +class TestHealthCheckCooldownIntegration: + """Test that health check failures trigger cooldown via _set_cooldown_deployments.""" + + def test_health_check_failure_increments_failed_calls(self): + """Health check failure should increment the failed_calls counter.""" + from litellm.router_utils.cooldown_handlers import ( + should_cooldown_based_on_allowed_fails_policy, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=3), + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout", model="gpt-4", llm_provider="openai" + ) + + # First call: should not cooldown (1 <= 3) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=timeout_exc, + ) + assert result is False + + # Check counter was incremented + current_fails = router.failed_calls.get_cache(key="deploy-1") + assert current_fails == 1 + + def test_health_check_failure_triggers_cooldown_at_threshold(self): + """After exceeding allowed_fails threshold, deployment should enter cooldown.""" + from litellm.router_utils.cooldown_handlers import ( + should_cooldown_based_on_allowed_fails_policy, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=2), + ) + + auth_exc = litellm.AuthenticationError( + message="Invalid key", model="gpt-4", llm_provider="openai" + ) + + # Fails 1 and 2: should not cooldown + for _ in range(2): + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=auth_exc, + ) + assert result is False + + # Fail 3: should trigger cooldown (3 > 2) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=auth_exc, + ) + assert result is True + + def test_health_check_failure_falls_back_to_allowed_fails(self): + """When policy has no matching field, fall back to generic allowed_fails.""" + from litellm.router_utils.cooldown_handlers import ( + should_cooldown_based_on_allowed_fails_policy, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=10), + allowed_fails=1, + ) + + # Use an exception that doesn't match TimeoutErrorAllowedFails + # InternalServerError is not checked by get_allowed_fails_from_policy + # so it will fall back to allowed_fails=1 + generic_exc = Exception("Some internal error") + + # Fail 1: should not cooldown (1 <= 1) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=generic_exc, + ) + assert result is False + + # Fail 2: should trigger cooldown (2 > 1) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=generic_exc, + ) + assert result is True + + def test_healthy_endpoints_do_not_trigger_cooldown(self): + """Healthy endpoints should not increment any failure counters.""" + from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments + + router = Router( + model_list=[_make_model("deploy-1")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=1), + enable_health_check_routing=True, + ) + + # Simulate healthy endpoint -- no exception, no cooldown call + healthy_endpoint = {"model_id": "deploy-1"} + # Should have no exception key + assert "exception" not in healthy_endpoint + + # Verify failed_calls counter is untouched + current_fails = router.failed_calls.get_cache(key="deploy-1") + assert current_fails is None + + def test_disable_cooldowns_prevents_health_check_cooldown(self): + """When disable_cooldowns=True, health check failures should not trigger cooldown.""" + from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0), + enable_health_check_routing=True, + disable_cooldowns=True, + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout", model="gpt-4", llm_provider="openai" + ) + + result = _set_cooldown_deployments( + litellm_router_instance=router, + original_exception=timeout_exc, + exception_status=500, + deployment="deploy-1", + time_to_cooldown=router.cooldown_time, + ) + assert result is False + + +class TestWriteHealthStateIntegration: + """Test _write_health_state_to_router_cache integrates with cooldown pipeline.""" + + def test_unhealthy_endpoint_triggers_set_cooldown(self): + """_write_health_state_to_router_cache should call _set_cooldown_deployments for unhealthy endpoints.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5), + enable_health_check_routing=True, + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout", model="", llm_provider="" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "timeout"}, + ] + healthy_endpoints = [ + {"model_id": "deploy-2"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=healthy_endpoints, + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, + ) + mock_cooldown.assert_called_once_with( + litellm_router_instance=router, + original_exception=timeout_exc, + exception_status=408, # Timeout has status_code 408 + deployment="deploy-1", + time_to_cooldown=router.cooldown_time, + ) + + def test_unhealthy_endpoint_without_exception_skips_cooldown(self): + """Unhealthy endpoints without an exception key should not trigger cooldown.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5), + enable_health_check_routing=True, + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "unknown failure"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + # no exceptions_by_model_id → cooldown should not fire + ) + mock_cooldown.assert_not_called() + + def test_unhealthy_endpoint_increments_failure_counter(self): + """Unhealthy endpoints should call increment_deployment_failures_for_current_minute.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=10), + enable_health_check_routing=True, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute" + ) as mock_increment: + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + mock_increment.assert_called_once_with( + litellm_router_instance=router, + deployment_id="deploy-1", + ) + + +class TestHealthCheckFilterBypassWithPolicy: + """ + When allowed_fails_policy is set, the binary health check filter should be + bypassed so cooldown is the sole routing exclusion mechanism. + """ + + def test_filter_bypassed_when_policy_set(self): + """Binary health check filter is a no-op when allowed_fails_policy is configured.""" + import time + + from litellm.caching.caching import DualCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), + enable_health_check_routing=True, + ) + + # Mark deploy-1 as unhealthy in the health state cache + cache = DualCache() + health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + health_cache.set_deployment_health_states( + { + "deploy-1": { + "is_healthy": False, + "timestamp": time.time(), + "reason": "test", + }, + } + ) + router.health_state_cache = health_cache + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + # Filter should pass all through because policy is set + result = router._filter_health_check_unhealthy_deployments(deployments) + assert ( + len(result) == 2 + ), "Binary filter should be bypassed when allowed_fails_policy is set" + + def test_filter_active_when_no_policy(self): + """Binary health check filter still works when no allowed_fails_policy is configured.""" + import time + + from litellm.caching.caching import DualCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + ) + + cache = DualCache() + health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + health_cache.set_deployment_health_states( + { + "deploy-1": { + "is_healthy": False, + "timestamp": time.time(), + "reason": "test", + }, + } + ) + router.health_state_cache = health_cache + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + result = router._filter_health_check_unhealthy_deployments(deployments) + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "deploy-2" + + @pytest.mark.asyncio + async def test_async_filter_bypassed_when_policy_set(self): + """Async version also bypasses when allowed_fails_policy is set.""" + import time + + from litellm.caching.caching import DualCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=2), + enable_health_check_routing=True, + ) + + cache = DualCache() + health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + health_cache.set_deployment_health_states( + { + "deploy-1": { + "is_healthy": False, + "timestamp": time.time(), + "reason": "test", + }, + } + ) + router.health_state_cache = health_cache + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + result = await router._async_filter_health_check_unhealthy_deployments( + deployments + ) + assert len(result) == 2 + + +class TestAllDeploymentsInCooldownSafetyNet: + """ + When enable_health_check_routing=True and ALL deployments enter cooldown, + the async routing path should bypass the cooldown filter and return all + deployments rather than blocking all traffic. + """ + + def test_raw_cooldown_filter_returns_empty_when_all_cooled(self): + """The raw _filter_cooldown_deployments has no safety net -- it returns empty.""" + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + ) + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + result = router._filter_cooldown_deployments( + healthy_deployments=deployments, + cooldown_deployments=["deploy-1", "deploy-2"], + ) + assert result == [] # raw filter has no safety net + + @pytest.mark.asyncio + async def test_async_routing_path_bypasses_all_cooldown(self): + """In the async routing path, all-in-cooldown with enable_health_check_routing + returns the full list instead of empty (safety net).""" + from unittest.mock import AsyncMock + + from litellm.router_utils.cooldown_handlers import ( + _async_get_cooldown_deployments, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0), + enable_health_check_routing=True, + ) + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + # Simulate all deployments in cooldown + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["deploy-1", "deploy-2"]), + ): + # The safety net in async_get_available_deployment should restore + # all deployments when the cooldown filter empties the list + _pre = deployments.copy() + filtered = router._filter_cooldown_deployments( + healthy_deployments=deployments, + cooldown_deployments=["deploy-1", "deploy-2"], + ) + # If filtered is empty and enable_health_check_routing is True, + # the routing path restores _pre_cooldown_deployments + if not filtered and router.enable_health_check_routing: + filtered = _pre + + assert ( + len(filtered) == 2 + ), "Safety net should return all deployments when all are in cooldown" + + +class TestHealthCheckIgnoreTransientErrors: + """ + When health_check_ignore_transient_errors=True, health check failures with + 429 or 408 status codes should NOT increment failure counters or trigger cooldown. + 401, 404, and 5xx errors should still be processed normally. + """ + + def test_429_skipped_when_flag_enabled(self): + """429 from health check does not trigger cooldown when flag is set.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + assert getattr(rate_exc, "status_code", None) == 429 + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + with patch( + "litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute" + ) as mock_increment: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + mock_cooldown.assert_not_called() + mock_increment.assert_not_called() + + def test_408_skipped_when_flag_enabled(self): + """408 from health check does not trigger cooldown when flag is set.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout exceeded", model="", llm_provider="" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "timeout"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, + ) + mock_cooldown.assert_not_called() + + def test_401_still_triggers_cooldown_when_flag_enabled(self): + """Auth errors (401) still trigger cooldown even when flag is set.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + auth_exc = litellm.AuthenticationError( + message="Invalid key", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "auth failed"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": auth_exc}, + ) + mock_cooldown.assert_called_once() + + def test_429_not_written_to_health_state_cache_when_flag_enabled(self): + """429 endpoint is excluded from health state cache when flag is set, + so the binary health check filter does not mark it as unhealthy.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1")], + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + + # Health state cache should have NO entry for deploy-1 + # (429 was ignored, not written as unhealthy) + unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids() + assert "deploy-1" not in unhealthy_ids + + def test_429_triggers_cooldown_when_flag_disabled(self): + """When flag is False (default), 429 still triggers cooldown.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=False, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + mock_cooldown.assert_called_once() + + +class TestSharedCacheTransientErrorFilter: + """ + When SharedHealthCheckManager returns cached results, exceptions_by_model_id + is always {}. The filter must fall back to the 'exception_status' field stored + on each endpoint dict so 429/408 endpoints are still excluded correctly. + """ + + def test_cached_429_excluded_via_exception_status_field(self): + """Cache-hit path: endpoint with exception_status=429 is excluded from health state.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + # Simulate a cache-hit endpoint: exception_status stored as int, no exceptions dict + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited", "exception_status": 429}, + ] + + with patch.object(proxy_module, "llm_router", router): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={}, + ) + + # deploy-1 should NOT be marked unhealthy (429 was filtered) + unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids() + assert "deploy-1" not in unhealthy_ids + + def test_cached_401_still_marked_unhealthy(self): + """Cache-hit path: endpoint with exception_status=401 is still written as unhealthy.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "auth failed", "exception_status": 401}, + ] + + with patch.object(proxy_module, "llm_router", router): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={}, + ) + + unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids() + assert "deploy-1" in unhealthy_ids diff --git a/tests/test_litellm/router_utils/test_router_health_check_routing.py b/tests/test_litellm/router_utils/test_router_health_check_routing.py index f40144b44c9..b87a39ac1de 100644 --- a/tests/test_litellm/router_utils/test_router_health_check_routing.py +++ b/tests/test_litellm/router_utils/test_router_health_check_routing.py @@ -50,6 +50,7 @@ class TestFilterHealthCheckUnhealthyDeployments: def __init__(self): self.enable_health_check_routing = enable self.health_state_cache = health_cache + self.allowed_fails_policy = None # Import the actual method and bind it from litellm.router import Router @@ -125,6 +126,7 @@ class TestAsyncFilterHealthCheckUnhealthyDeployments: def __init__(self): self.enable_health_check_routing = enable self.health_state_cache = health_cache + self.allowed_fails_policy = None fake = FakeRouter() fake._async_filter_health_check_unhealthy_deployments = ( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3df10901492..262dce439c0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5,7 +5,6 @@ import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") @@ -13,7 +12,6 @@ sys.path.insert( import litellm -from litellm.router_utils.fallback_event_handlers import run_async_fallback def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -127,7 +125,7 @@ async def test_async_router_acreate_file(): """ Write to all deployments of a model """ - from unittest.mock import MagicMock, call, patch + from unittest.mock import MagicMock, patch router = litellm.Router( model_list=[ @@ -747,7 +745,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): """ Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import patch router = litellm.Router( model_list=[ @@ -1134,10 +1132,9 @@ def test_get_model_access_groups_cache_invalidation_upsert_deployment(): @pytest.mark.asyncio async def test_acompletion_streaming_iterator(): """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" - from unittest.mock import AsyncMock, MagicMock + from unittest.mock import MagicMock from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import ModelResponseStream # Helper class for creating async iterators class AsyncIterator: @@ -2847,3 +2844,286 @@ def test_combine_fallback_usage(): assert chunk.usage.prompt_tokens == 10 assert chunk.usage.completion_tokens == 5 assert chunk.usage.total_tokens == 15 + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback(): + """ + Test that fallback works correctly for team-scoped models. + + When a team-scoped model fails and the fallback model is also team-scoped, + the router should find the fallback deployment by matching team_public_model_name. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "team-a-fallback-internal", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "fallback success from team-a", + }, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "fallback-model", + }, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + assert response is not None + assert response.choices[0].message.content == "fallback success from team-a" + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback_to_global(): + """ + Test that a team-scoped model can fall back to a global (non-team) model. + + Global models (no team_id on deployment) should be accessible as fallback + targets for team-scoped requests. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "global-fallback", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "global fallback success", + }, + }, + ], + fallbacks=[{"primary-model": ["global-fallback"]}], + ) + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + assert response is not None + assert response.choices[0].message.content == "global fallback success" + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback_cross_team_blocked(): + """ + Test that cross-team fallback is correctly blocked. + + When team-a's model fails and the fallback target is scoped to team-b, + the router should NOT use it (team isolation). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "team-b-fallback-internal", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "team-b response - should not reach here", + }, + "model_info": { + "team_id": "team-b", + "team_public_model_name": "fallback-model", + }, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + with pytest.raises(Exception): + await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + + +def test_get_all_deployments_with_team_id(): + """ + Test that _get_all_deployments with team_id can find deployments + by team_public_model_name when the model_name is not in the index. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "internal-team-deployment", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": { + "team_id": "team-x", + "team_public_model_name": "gpt-4", + }, + }, + ], + ) + + # Without team_id: "gpt-4" is not in the model_name index (internal name is different) + deployments = router._get_all_deployments(model_name="gpt-4") + assert len(deployments) == 0 + + # With correct team_id: should find via O(n) scan matching team_public_model_name + deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-x") + assert len(deployments) == 1 + assert deployments[0]["model_name"] == "internal-team-deployment" + + # With wrong team_id: should find nothing + deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-y") + assert len(deployments) == 0 + + +def test_multiregion_team_deployments_unique_model_names(): + """ + Simulates athenahealth's exact setup: unique model_names per deployment, + same team_public_model_name, multiple regions. + + Verifies that _get_all_deployments returns ALL regional deployments + for a team when queried by team_public_model_name. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "metis-claude-us-east-1", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "aws_region_name": "us-east-1", + "api_key": "fake", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + { + "model_name": "metis-claude-us-west-2", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "aws_region_name": "us-west-2", + "api_key": "fake", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + ], + ) + + # "claude-sonnet" is NOT in the model_name index + assert "claude-sonnet" not in router.model_names + + # Without team_id: returns nothing (no model_name="claude-sonnet" in index, no O(n) scan) + deployments = router._get_all_deployments(model_name="claude-sonnet") + assert len(deployments) == 0 + + # With team_id: O(n) scan finds BOTH regional deployments + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert len(deployments) == 2 + deployment_names = {d["model_name"] for d in deployments} + assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} + + # Each deployment has a unique ID (critical for cooldown/retry to work) + deployment_ids = {d["model_info"]["id"] for d in deployments} + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" + + # Wrong team: returns nothing + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="other-team" + ) + assert len(deployments) == 0 + + +@pytest.mark.asyncio +async def test_multiregion_team_failover_between_regions(): + """ + Simulates athenahealth's multiregion failover scenario: + - Two Bedrock deployments (us-east-1 and us-west-2) with unique model_names + - Same team_public_model_name ("claude-sonnet") + - Primary region fails → router should failover to second region + + This is the exact scenario Sean Glover from athenahealth will demonstrate. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "metis-claude-us-east-1", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "api_key": "fake", + "mock_response": "response from us-east-1", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + { + "model_name": "metis-claude-us-west-2", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "api_key": "fake", + "mock_response": "response from us-west-2", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + ], + num_retries=1, + ) + + # Verify the router finds both deployments for the team + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert len(deployments) == 2, ( + "Router must find both regional deployments by team_public_model_name" + ) + + # Make a normal request — should succeed from one of the regions + response = await router.acompletion( + model="claude-sonnet", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "metis-team"}, + ) + assert response is not None + assert response.choices[0].message.content in [ + "response from us-east-1", + "response from us-west-2", + ] diff --git a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py index 21fecc015a3..c7d98548876 100644 --- a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py +++ b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py @@ -21,6 +21,7 @@ def test_pipeline_step_defaults(): step = PipelineStep(guardrail="my-guard") assert step.on_fail == "block" assert step.on_pass == "allow" + assert step.on_error is None assert step.pass_data is False assert step.modify_response_message is None @@ -33,9 +34,10 @@ def test_pipeline_step_valid_actions(): def test_pipeline_step_all_action_types(): for action in ("allow", "block", "next", "modify_response"): - step = PipelineStep(guardrail="g", on_fail=action, on_pass=action) + step = PipelineStep(guardrail="g", on_fail=action, on_pass=action, on_error=action) assert step.on_fail == action assert step.on_pass == action + assert step.on_error == action def test_pipeline_step_invalid_action_rejected(): @@ -48,6 +50,16 @@ def test_pipeline_step_invalid_on_pass_rejected(): PipelineStep(guardrail="my-guard", on_pass="skip") +def test_pipeline_step_on_error_valid(): + step = PipelineStep(guardrail="g", on_error="next", on_fail="block", on_pass="allow") + assert step.on_error == "next" + + +def test_pipeline_step_invalid_on_error_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_error="invalid") + + def test_pipeline_requires_at_least_one_step(): with pytest.raises(ValidationError): GuardrailPipeline(mode="pre_call", steps=[]) diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index b1768d5b81c..ba639594999 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -155,6 +155,14 @@ const FailIcon: React.FC = () => ( ); +const ApiFailureIcon: React.FC = () => ( + + + + + +); + // ───────────────────────────────────────────────────────────────────────────── // Connector // ───────────────────────────────────────────────────────────────────────────── @@ -349,6 +357,41 @@ const StepCard: React.FC = ({
)} + + {/* ON API FAILURE (technical / provider outage) — optional; defaults to ON FAIL */} +
+
+ + ON API FAILURE +
+ + + + + + +
+ + +
+ {/* Left panel: Available Tools */} +
+
+ Available Tools +
+ setServerSearch(e.target.value)} + className="mb-2" + allowClear + /> +
+ {filteredServers.length === 0 ? ( + {mcpServers.length === 0 ? "No MCP servers configured" : "No servers match your search"} + ) : ( + filteredServers.map((server) => ( + + )) + )} +
+
+ + {/* Divider */} +
+ + {/* Right panel: Your Toolset */} +
+ + Your Toolset{" "} + ({selectedTools.length} tools) + +
+ {selectedTools.length === 0 ? ( + No tools added yet + ) : ( + selectedTools.map((tool, idx) => ( + + )) + )} +
+
+
+ +
+ + +
+ + ); +} + +function toolsetColumns( + isAdmin: boolean, + onEdit: (t: MCPToolset) => void, + onDelete: (id: string) => void, + proxyBaseUrl: string, +): ColumnDef[] { + return [ + { + header: "Toolset ID", + accessorKey: "toolset_id", + cell: ({ row }) => ( + + {row.original.toolset_id.slice(0, 8)}… + + ), + }, + { + header: "Name", + accessorKey: "toolset_name", + cell: ({ row }) => { + const url = `${proxyBaseUrl}/toolset/${row.original.toolset_name}/mcp`; + return ( +
+
+ + {row.original.toolset_name} +
+ +
+ ); + }, + }, + { + header: "Description", + accessorKey: "description", + cell: ({ row }) => ( + {row.original.description || "—"} + ), + }, + { + header: "Tools", + accessorKey: "tools", + cell: ({ row }) => { + const tools = row.original.tools; + return ( +
+ {tools.slice(0, 4).map((t, i) => ( + + {t.tool_name} + + ))} + {tools.length > 4 && ( + +{tools.length - 4} more + )} +
+ ); + }, + }, + { + header: "Created", + accessorKey: "created_at", + cell: ({ row }) => ( + + {row.original.created_at ? new Date(row.original.created_at).toLocaleDateString() : "—"} + + ), + }, + ...(isAdmin ? [{ + header: "", + id: "actions", + cell: ({ row }: { row: { original: MCPToolset } }) => ( +
+ + +
+ ), + } as ColumnDef] : []), + ]; +} + +function ToolsetUsageGuide() { + const [copied, setCopied] = useState(false); + const proxyBaseUrl = getProxyBaseUrl(); + + const snippet = `{ + "mcpServers": { + "my-toolset": { + "url": "${proxyBaseUrl}/toolset//mcp", + "headers": { "x-litellm-api-key": "Bearer " } + } + } +}`; + + const copy = async () => { + try { + await navigator.clipboard.writeText(snippet); + setCopied(true); + setTimeout(() => setCopied(false), 1500); + } catch { + // ignore + } + }; + + return ( +
+

How toolsets work

+

+ Create a toolset, assign it to a key via API Keys → Edit Key → MCP Servers, then point your MCP client at the toolset URL. The client only sees the tools you picked. +

+
Claude Code / Cursor config
+
+ 
+          {snippet}
+
+ +
+
+ ); +} + +export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { + const queryClient = useQueryClient(); + const { data: toolsets = [], isLoading } = useMCPToolsets(); + const [createOpen, setCreateOpen] = useState(false); + const [editToolset, setEditToolset] = useState(null); + const [deleteId, setDeleteId] = useState(null); + const [deleting, setDeleting] = useState(false); + + const isAdmin = userRole === "Admin" || userRole === "proxy_admin"; + + const handleCreate = async (name: string, description: string | undefined, tools: MCPToolsetTool[]) => { + if (!accessToken) return; + await createMCPToolset(accessToken, { toolset_name: name, description, tools }); + message.success("Toolset created"); + queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); + }; + + const handleUpdate = async (name: string, description: string | undefined, tools: MCPToolsetTool[]) => { + if (!accessToken || !editToolset) return; + await updateMCPToolset(accessToken, { toolset_id: editToolset.toolset_id, toolset_name: name, description, tools }); + message.success("Toolset updated"); + queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); + setEditToolset(null); + }; + + const handleDelete = async () => { + if (!accessToken || !deleteId) return; + setDeleting(true); + try { + await deleteMCPToolset(accessToken, deleteId); + message.success("Toolset deleted"); + queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); + setDeleteId(null); + } finally { + setDeleting(false); + } + }; + + const proxyBaseUrl = getProxyBaseUrl(); + const columns = toolsetColumns(isAdmin, setEditToolset, setDeleteId, proxyBaseUrl); + + return ( +
+
+
+ MCP Toolsets + + Curated collections of tools from one or more MCP servers. Assign toolsets to keys and teams via the MCP permissions dropdown. + +
+ {isAdmin && ( + + )} +
+ + + +
} + getRowCanExpand={() => false} + isLoading={isLoading} + noDataMessage="No toolsets yet. Click 'New Toolset' to create one." + loadingMessage="Loading toolsets..." + enableSorting={true} + /> + + setCreateOpen(false)} + onSave={handleCreate} + accessToken={accessToken} + /> + + {editToolset && ( + setEditToolset(null)} + onSave={handleUpdate} + accessToken={accessToken} + initialToolset={editToolset} + /> + )} + + setDeleteId(null)} + onOk={handleDelete} + okText="Delete" + okButtonProps={{ danger: true, loading: deleting }} + title="Delete Toolset" + > +

Are you sure you want to delete this toolset? Keys and teams using it will lose access to the scoped tools.

+ +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx index 4d67456b78a..fc1d941c177 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx @@ -102,7 +102,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); const queryClient = createQueryClient(); - const { getByText } = render( + const { getByText, getAllByText } = render( , @@ -121,8 +121,8 @@ describe("MCPServers", () => { // Verify the mocked server data is rendered in the table expect(getByText("Test Server 1")).toBeInTheDocument(); expect(getByText("Test Server 2")).toBeInTheDocument(); - expect(getByText("test-server-1")).toBeInTheDocument(); - expect(getByText("test-server-2")).toBeInTheDocument(); + expect(getAllByText("test-server-1").length).toBeGreaterThan(0); + expect(getAllByText("test-server-2").length).toBeGreaterThan(0); // Verify the API was called // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 34eeb6b1e86..f0fed7c7fec 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -9,6 +9,7 @@ import { useMCPServerHealth } from "../../app/(dashboard)/hooks/mcpServers/useMC import NotificationsManager from "../molecules/notifications_manager"; import { deleteMCPServer } from "../networking"; import { MCPSubmissionsTab } from "./MCPSubmissionsTab"; +import { MCPToolsetsTab } from "./MCPToolsetsTab"; import { DataTable } from "../view_logs/table"; import CreateMCPServer from "./create_mcp_server"; import MCPConnect from "./mcp_connect"; @@ -348,6 +349,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID })
All Servers + Toolsets Connect Semantic Filter Network Settings @@ -426,6 +428,9 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID })
)} + + + diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 67526ca2b07..acf8310c803 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -231,6 +231,20 @@ export interface MCPServerProps { userID: string | null; } +export interface MCPToolsetTool { + server_id: string; + tool_name: string; +} + +export interface MCPToolset { + toolset_id: string; + toolset_name: string; + description?: string; + tools: MCPToolsetTool[]; + created_at?: string; + created_by?: string; +} + // Discoverable MCP server from the curated registry export interface DiscoverableMCPServer { name: string; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 7f8dea4cbfe..28f8d308de7 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6675,6 +6675,99 @@ export const deleteMCPServer = async (accessToken: string, serverId: string) => } }; +export const fetchMCPToolsets = async (accessToken: string): Promise => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset`; + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + return await response.json(); + } catch (error) { + console.error("Failed to fetch MCP toolsets:", error); + throw error; + } +}; + +export const createMCPToolset = async (accessToken: string, formValues: Record) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset`; + const response = await fetch(url, { + method: HTTP_REQUEST.POST, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(formValues), + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + return await response.json(); + } catch (error) { + console.error("Failed to create MCP toolset:", error); + throw error; + } +}; + +export const updateMCPToolset = async (accessToken: string, formValues: Record) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset`; + const response = await fetch(url, { + method: HTTP_REQUEST.PUT, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(formValues), + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + return await response.json(); + } catch (error) { + console.error("Failed to update MCP toolset:", error); + throw error; + } +}; + +export const deleteMCPToolset = async (accessToken: string, toolsetId: string) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset/${toolsetId}`; + const response = await fetch(url, { + method: HTTP_REQUEST.DELETE, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + } catch (error) { + console.error("Failed to delete MCP toolset:", error); + throw error; + } +}; + export const registerMCPServer = async (accessToken: string, formValues: Record) => { try { const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/server/register`; diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index ac55a57c7c3..685467e1d3e 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -9,6 +9,7 @@ interface ObjectPermission { mcp_servers: string[]; mcp_access_groups?: string[]; mcp_tool_permissions?: Record; + mcp_toolsets?: string[]; vector_stores: string[]; agents?: string[]; agent_access_groups?: string[]; @@ -31,17 +32,19 @@ export function ObjectPermissionsView({ const mcpServers = objectPermission?.mcp_servers || []; const mcpAccessGroups = objectPermission?.mcp_access_groups || []; const mcpToolPermissions = objectPermission?.mcp_tool_permissions || {}; + const mcpToolsets = objectPermission?.mcp_toolsets || []; const agents = objectPermission?.agents || []; const agentAccessGroups = objectPermission?.agent_access_groups || []; const content = (
- { ); // Verify empty state message - expect(screen.getByText("No MCP servers or access groups configured")).toBeInTheDocument(); + expect(screen.getByText("No MCP servers, access groups, or toolsets configured")).toBeInTheDocument(); // Verify count badge shows 0 expect(screen.getByText("0")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index 962fc00670a..649e8efdf4a 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -2,25 +2,28 @@ import React, { useState, useEffect } from "react"; import { Text, Badge } from "@tremor/react"; import { ServerIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; import { Tooltip } from "antd"; -import { fetchMCPServers } from "../networking"; -import { MCPServer } from "../mcp_tools/types"; +import { fetchMCPServers, fetchMCPToolsets } from "../networking"; +import { MCPServer, MCPToolset } from "../mcp_tools/types"; interface MCPServerPermissionsProps { mcpServers: string[]; mcpAccessGroups?: string[]; mcpToolPermissions?: Record; + mcpToolsets?: string[]; accessToken?: string | null; } -export function MCPServerPermissions({ - mcpServers, - mcpAccessGroups = [], +export function MCPServerPermissions({ + mcpServers, + mcpAccessGroups = [], mcpToolPermissions = {}, - accessToken + mcpToolsets = [], + accessToken }: MCPServerPermissionsProps) { const [mcpServerDetails, setMCPServerDetails] = useState([]); - const [accessGroupNames, setAccessGroupNames] = useState([]); + const [toolsetDetails, setToolsetDetails] = useState([]); const [expandedServers, setExpandedServers] = useState>(new Set()); + const [expandedToolsets, setExpandedToolsets] = useState>(new Set()); const toggleServerExpansion = (serverId: string) => { setExpandedServers((prev) => { @@ -34,6 +37,18 @@ export function MCPServerPermissions({ }); }; + const toggleToolsetExpansion = (toolsetId: string) => { + setExpandedToolsets((prev) => { + const newSet = new Set(prev); + if (newSet.has(toolsetId)) { + newSet.delete(toolsetId); + } else { + newSet.add(toolsetId); + } + return newSet; + }); + }; + // Fetch MCP server details when component mounts useEffect(() => { const fetchMCPServerDetails = async () => { @@ -53,20 +68,23 @@ export function MCPServerPermissions({ fetchMCPServerDetails(); }, [accessToken, mcpServers.length]); - // Fetch MCP access group names + // Fetch toolset details useEffect(() => { - const fetchGroups = async () => { - if (accessToken && mcpAccessGroups.length > 0) { + const fetchToolsets = async () => { + if (accessToken && mcpToolsets.length > 0) { try { - const groups = await import("../networking").then((m) => m.fetchMCPAccessGroups(accessToken)); - setAccessGroupNames(Array.isArray(groups) ? groups : groups.data || []); + const all = await fetchMCPToolsets(accessToken); + const filtered = Array.isArray(all) + ? all.filter((t: MCPToolset) => mcpToolsets.includes(t.toolset_id)) + : []; + setToolsetDetails(filtered); } catch (error) { - console.error("Error fetching MCP access groups:", error); + console.error("Error fetching toolsets:", error); } } }; - fetchGroups(); - }, [accessToken, mcpAccessGroups.length]); + fetchToolsets(); + }, [accessToken, mcpToolsets.length]); // Function to get display name for MCP server const getMCPServerDisplayName = (serverId: string) => { @@ -78,17 +96,12 @@ export function MCPServerPermissions({ return serverId; }; - // Function to get display name for access group - const getAccessGroupDisplayName = (group: string) => { - return group; - }; - // Merge servers and access groups into one list const mergedItems = [ ...mcpServers.map((server) => ({ type: "server", value: server })), ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })), ]; - const totalCount = mergedItems.length; + const totalCount = mergedItems.length + mcpToolsets.length; return (
@@ -99,21 +112,21 @@ export function MCPServerPermissions({ {totalCount}
- + {totalCount > 0 ? (
{mergedItems.map((item, index) => { const toolsForServer = item.type === "server" ? mcpToolPermissions[item.value] : undefined; const hasToolRestrictions = toolsForServer && toolsForServer.length > 0; const isExpanded = expandedServers.has(item.value); - + return (
-
hasToolRestrictions && toggleServerExpansion(item.value)} className={`flex items-center gap-3 py-2 px-3 rounded-lg border border-gray-200 transition-all ${ - hasToolRestrictions - ? 'cursor-pointer hover:bg-gray-50 hover:border-gray-300' + hasToolRestrictions + ? 'cursor-pointer hover:bg-gray-50 hover:border-gray-300' : 'bg-white' }`} > @@ -128,14 +141,14 @@ export function MCPServerPermissions({ ) : (
- {getAccessGroupDisplayName(item.value)} + {item.value} Group
)}
- + {hasToolRestrictions && (
{toolsForServer.length} @@ -148,7 +161,7 @@ export function MCPServerPermissions({
)}
- + {/* Show tool permissions if expanded */} {hasToolRestrictions && isExpanded && (
@@ -167,11 +180,66 @@ export function MCPServerPermissions({
); })} + + {/* Toolsets section */} + {mcpToolsets.length > 0 && mcpToolsets.map((toolsetId, index) => { + const detail = toolsetDetails.find((t) => t.toolset_id === toolsetId); + const isExpanded = expandedToolsets.has(toolsetId); + const toolCount = detail?.tools.length ?? 0; + + return ( +
+
toolCount > 0 && toggleToolsetExpansion(toolsetId)} + className={`flex items-center gap-3 py-2 px-3 rounded-lg border border-purple-200 transition-all ${ + toolCount > 0 ? 'cursor-pointer hover:bg-purple-50 hover:border-purple-300' : 'bg-white' + }`} + > +
+ + + {detail?.toolset_name ?? toolsetId} + + + Toolset + +
+ {toolCount > 0 && ( +
+ {toolCount} + {toolCount === 1 ? "tool" : "tools"} + {isExpanded ? ( + + ) : ( + + )} +
+ )} +
+ + {toolCount > 0 && isExpanded && detail && ( +
+
+ {detail.tools.map((tool, toolIndex) => ( + + {tool.server_id.slice(0, 6)}… + {tool.tool_name} + + ))} +
+
+ )} +
+ ); + })}
) : (
- No MCP servers or access groups configured + No MCP servers, access groups, or toolsets configured
)}
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index e80c2b2f742..2bfe08efae8 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -34,7 +34,8 @@ import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/M import { MCPServer } from "../../mcp_tools/types"; import { ByokCredentialModal } from "../../mcp_tools/ByokCredentialModal"; import NotificationsManager from "../../molecules/notifications_manager"; -import { callMCPTool, fetchMCPServers, listMCPTools } from "../../networking"; +import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "../../networking"; +import { MCPToolset } from "../../mcp_tools/types"; import TagSelector from "../../tag_management/TagSelector"; import VectorStoreSelector from "../../vector_store_management/VectorStoreSelector"; import { makeA2ASendMessageRequest } from "../llm_calls/a2a_send_message"; @@ -111,6 +112,8 @@ const ChatUI: React.FC = ({ fixedModel, }) => { const [mcpServers, setMCPServers] = useState([]); + const [mcpToolsets, setMCPToolsets] = useState([]); + const [isToolsetsInfoModalVisible, setIsToolsetsInfoModalVisible] = useState(false); const [byokModalServer, setByokModalServer] = useState(null); const [selectedMCPServers, setSelectedMCPServers] = useState(() => { const saved = sessionStorage.getItem("selectedMCPServers"); @@ -257,15 +260,19 @@ const ChatUI: React.FC = ({ const chatEndRef = useRef(null); - // Fetch MCP servers + // Fetch MCP servers and toolsets const loadMCPServers = async () => { const userApiKey = apiKeySource === "session" ? accessToken : apiKey; if (!userApiKey) return; setIsLoadingMCPServers(true); try { - const servers = await fetchMCPServers(userApiKey); + const [servers, toolsets] = await Promise.all([ + fetchMCPServers(userApiKey), + fetchMCPToolsets(userApiKey).catch(() => []), + ]); setMCPServers(Array.isArray(servers) ? servers : servers.data || []); + setMCPToolsets(Array.isArray(toolsets) ? toolsets : []); } catch (error) { console.error("Error fetching MCP servers:", error); } finally { @@ -416,17 +423,29 @@ const ChatUI: React.FC = ({ loadMCPServers(); }, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]); - // Load tools when MCP direct mode has a server selected + // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { if ( endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && - selectedMCPServers[0] !== "__all__" && - !serverToolsMap[selectedMCPServers[0]] + selectedMCPServers[0] !== "__all__" ) { - loadServerTools(selectedMCPServers[0]); + const selected = selectedMCPServers[0]; + if (selected.startsWith("toolset:")) { + // For a toolset, load tools for each server in it + const toolsetId = selected.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + if (toolset) { + const uniqueServerIds = [...new Set(toolset.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + if (!serverToolsMap[sid]) loadServerTools(sid); + }); + } + } else if (!serverToolsMap[selected]) { + loadServerTools(selected); + } } - }, [endpointType, selectedMCPServers, serverToolsMap]); + }, [endpointType, selectedMCPServers, serverToolsMap, mcpToolsets]); // Fetch agents when A2A endpoint is selected useEffect(() => { @@ -572,19 +591,34 @@ const ChatUI: React.FC = ({ // For MCP direct mode, require server and tool selection, and get form values early let mcpToolArguments: Record = {}; if (endpointType === EndpointType.MCP) { - const mcpServerId = + const rawSelected = selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" ? selectedMCPServers[0] : null; - if (!mcpServerId) { + if (!rawSelected) { NotificationsManager.fromBackend("Please select an MCP server to test"); return; } + // Resolve the real server ID (toolsets use toolset: prefix) + const mcpServerId = rawSelected.startsWith("toolset:") ? rawSelected : rawSelected; if (!selectedMCPDirectTool) { NotificationsManager.fromBackend("Please select an MCP tool to call"); return; } - const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find( + // For toolsets, find the tool in the servers that back this toolset + const toolsetForSelected = rawSelected.startsWith("toolset:") + ? mcpToolsets.find((t) => t.toolset_id === rawSelected.slice("toolset:".length)) + : null; + let searchPool: any[] = []; + if (toolsetForSelected) { + const uniqueServerIds = [...new Set(toolsetForSelected.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + searchPool = searchPool.concat(serverToolsMap[sid] || []); + }); + } else { + searchPool = serverToolsMap[rawSelected] || []; + } + const mcpTool = searchPool.find( (t: any) => t.name === selectedMCPDirectTool, ); if (!mcpTool) { @@ -742,6 +776,7 @@ const ChatUI: React.FC = ({ mcpServerToolRestrictions, handleMCPEvent, mockTestFallbacks, + mcpToolsets, ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -822,6 +857,7 @@ const ChatUI: React.FC = ({ customProxyBaseUrl || undefined, mcpServers, mcpServerToolRestrictions, + mcpToolsets, ); } else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) { const apiChatHistory = [ @@ -879,14 +915,22 @@ const ChatUI: React.FC = ({ // Handle MCP direct tool calls (no chat completions) if (endpointType === EndpointType.MCP) { - const mcpServerId = + const rawSelected = selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" ? selectedMCPServers[0] : null; - if (mcpServerId && selectedMCPDirectTool) { + // For toolsets, resolve the real server_id from the toolset's tool list + let resolvedServerId = rawSelected; + if (rawSelected?.startsWith("toolset:")) { + const toolsetId = rawSelected.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + const toolEntry = toolset?.tools.find((t) => t.tool_name === selectedMCPDirectTool); + resolvedServerId = toolEntry?.server_id ?? rawSelected; + } + if (resolvedServerId && !resolvedServerId.startsWith("toolset:") && selectedMCPDirectTool) { const result = await callMCPTool( effectiveApiKey, - mcpServerId, + resolvedServerId, selectedMCPDirectTool, mcpToolArguments, selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined, @@ -1301,11 +1345,14 @@ const ChatUI: React.FC = ({ className="ml-1" title={ endpointType === EndpointType.MCP - ? "Select an MCP server to test tools directly." - : "Select MCP servers to use in your conversation." + ? "Select an MCP server or toolset to test tools directly." + : "Select MCP servers or toolsets to use in your conversation." } > - + setIsToolsetsInfoModalVisible(true)} + /> {/* MCP Tool selector - only for MCP direct mode */} {endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && - selectedMCPServers[0] !== "__all__" && ( -
- Select Tool - setSelectedMCPDirectTool(value)} + options={toolOptions} + allowClear + className="rounded-md" + /> +
+ ); + })()} {/* Tool restrictions UI (optional) - hidden for MCP direct mode */} {selectedMCPServers.length > 0 && @@ -1924,7 +2037,21 @@ const ChatUI: React.FC = ({ selectedMCPDirectTool ? (
{(() => { - const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find( + const rawSel = selectedMCPServers[0]; + let toolPool: any[] = []; + if (rawSel.startsWith("toolset:")) { + const toolsetId = rawSel.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + if (toolset) { + const uniqueServerIds = [...new Set(toolset.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + toolPool = toolPool.concat(serverToolsMap[sid] || []); + }); + } + } else { + toolPool = serverToolsMap[rawSel] || []; + } + const mcpTool = toolPool.find( (t: any) => t.name === selectedMCPDirectTool, ); return mcpTool ? ( @@ -2070,6 +2197,47 @@ const ChatUI: React.FC = ({ accessToken={accessToken || ""} /> )} + + {/* Toolsets info modal */} + setIsToolsetsInfoModalVisible(false)} + footer={[ + , + ]} + width={600} + > +
+

+ Toolsets are named collections of specific tools from one or more MCP servers. + Instead of exposing all tools from a server, a toolset gives an agent exactly the tools it needs. +

+
+

How to use a toolset:

+
    +
  1. Select a Toolset (purple badge) from the MCP Servers dropdown.
    2. +
  2. The tool picker will show only the tools included in that toolset.
    3. +
  3. Select a tool and fill in its parameters, then send.
    4. +
  4. The tool call is routed to the correct underlying MCP server automatically.
    5. +
+
+
+

+ Example: A "GitHub Read-only" toolset might include only list_repos and get_file from a GitHub MCP server — preventing agents from making writes. +

+
+
+

Creating toolsets:

+

+ Admins can create and manage toolsets from the MCP page → Toolsets tab. + Toolsets can then be assigned to keys and teams to scope their tool access. +

+
+
+
); }; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index 3197c9409ce..b2281eaadce 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -3,7 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; -import { MCPServer, type MCPEvent } from "../../mcp_tools/types"; +import { MCPServer, MCPToolset, type MCPEvent } from "../../mcp_tools/types"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], @@ -30,6 +30,7 @@ export async function makeOpenAIChatCompletionRequest( mcpServerToolRestrictions?: Record, onMCPEvent?: (event: MCPEvent) => void, mockTestFallbacks?: boolean, + mcpToolsets?: MCPToolset[], ) { // base url should be the current base_url const isLocal = process.env.NODE_ENV === "development"; @@ -82,19 +83,31 @@ export async function makeOpenAIChatCompletionRequest( require_approval: "never", }); } else { - // Individual servers selected - create one entry per server + // Individual servers/toolsets selected - create one entry per item selectedMCPServers.forEach((serverId) => { - const server = mcpServers?.find((s) => s.server_id === serverId); - const serverName = server?.alias || server?.server_name || serverId; - const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + if (serverId.startsWith("toolset:")) { + const toolsetId = serverId.slice("toolset:".length); + const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); + const toolsetName = toolset?.toolset_name || toolsetId; + tools.push({ + type: "mcp", + server_label: toolsetName, + server_url: `litellm_proxy/mcp/${encodeURIComponent(toolsetName)}`, + require_approval: "never", + }); + } else { + const server = mcpServers?.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; - tools.push({ - type: "mcp", - server_label: "litellm", - server_url: `litellm_proxy/mcp/${serverName}`, - require_approval: "never", - ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), - }); + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: `litellm_proxy/mcp/${serverName}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + } }); } } diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx index 48d0efca6ee..4e88a356cf3 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx @@ -4,7 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; import type { MCPEvent } from "../../mcp_tools/types"; -import { MCPServer } from "../../mcp_tools/types"; +import { MCPServer, MCPToolset } from "../../mcp_tools/types"; import { CodeInterpreterResult, CodeInterpreterState, @@ -37,6 +37,7 @@ export async function makeOpenAIResponsesRequest( customBaseUrl?: string, mcpServers?: MCPServer[], mcpServerToolRestrictions?: Record, + mcpToolsets?: MCPToolset[], ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -102,21 +103,34 @@ export async function makeOpenAIResponsesRequest( require_approval: "never", }); } else { - // Individual servers selected - create one entry per server + // Individual servers/toolsets selected - create one entry per item selectedMCPServers.forEach((serverId) => { - const server = mcpServers?.find((s) => s.server_id === serverId); - // Use server_name for both routing and labelling. server_name is the - // unique registered identifier; aliases can collide across servers. - const routeName = server?.server_name || serverId; - const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + if (serverId.startsWith("toolset:")) { + // Toolset: same /{name}/mcp pattern as individual servers + const toolsetId = serverId.slice("toolset:".length); + const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); + const toolsetName = toolset?.toolset_name || toolsetId; + tools.push({ + type: "mcp", + server_label: toolsetName, + server_url: `${proxyBaseUrl}/mcp/${encodeURIComponent(toolsetName)}`, + require_approval: "never", + }); + } else { + const server = mcpServers?.find((s) => s.server_id === serverId); + // Use server_name for both routing and labelling. server_name is the + // unique registered identifier; aliases can collide across servers. + const routeName = server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; - tools.push({ - type: "mcp", - server_label: routeName, // unique per request — collisions cause silent tool-routing failures - server_url: `${proxyBaseUrl}/mcp/${encodeURIComponent(routeName)}`, - require_approval: "never", - ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), - }); + tools.push({ + type: "mcp", + server_label: routeName, // unique per request — collisions cause silent tool-routing failures + server_url: `${proxyBaseUrl}/mcp/${encodeURIComponent(routeName)}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + } }); } } diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index abc967bb984..e59af82e97b 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -105,6 +105,7 @@ export interface TeamData { mcp_servers: string[]; mcp_access_groups?: string[]; mcp_tool_permissions?: Record; + mcp_toolsets?: string[]; vector_stores: string[]; agents?: string[]; agent_access_groups?: string[]; @@ -516,9 +517,10 @@ const TeamInfoView: React.FC = ({ } // Handle object_permission updates - const { servers, accessGroups } = values.mcp_servers_and_groups || { + const { servers, accessGroups, toolsets } = values.mcp_servers_and_groups || { servers: [], accessGroups: [], + toolsets: [], }; const serverIds = new Set(servers || []); const mcpToolPermissions = Object.fromEntries( @@ -535,6 +537,9 @@ const TeamInfoView: React.FC = ({ if (mcpToolPermissions) { updateData.object_permission.mcp_tool_permissions = mcpToolPermissions; } + if (toolsets) { + updateData.object_permission.mcp_toolsets = toolsets; + } delete values.mcp_servers_and_groups; delete values.mcp_tool_permissions; @@ -838,6 +843,7 @@ const TeamInfoView: React.FC = ({ mcp_servers_and_groups: { servers: info.object_permission?.mcp_servers || [], accessGroups: info.object_permission?.mcp_access_groups || [], + toolsets: info.object_permission?.mcp_toolsets || [], }, mcp_tool_permissions: info.object_permission?.mcp_tool_permissions || {}, agents_and_groups: { diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 5d00ab3d0b9..24e3e18b93c 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -160,11 +160,12 @@ export default function KeyInfoView({ } if (formValues.mcp_servers_and_groups !== undefined) { - const { servers, accessGroups } = formValues.mcp_servers_and_groups || { servers: [], accessGroups: [] }; + const { servers, accessGroups, toolsets } = formValues.mcp_servers_and_groups || { servers: [], accessGroups: [], toolsets: [] }; formValues.object_permission = { ...currentKeyData.object_permission, mcp_servers: servers || [], mcp_access_groups: accessGroups || [], + mcp_toolsets: toolsets || [], }; // Remove mcp_servers_and_groups from the top level as it should be in object_permission delete formValues.mcp_servers_and_groups; From a36fe70fde0bc6dddfb9504ba623d220b01e6713 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 16:32:55 -0700 Subject: [PATCH 45/55] test(ui): fix guardrails.test.tsx after antd Tabs migration The test was hitting "No QueryClient set" because TeamGuardrailsTab (which pulls in useRegisterGuardrail) was not mocked alongside the other tab children. Added a mock. Also: the "+ Add New Guardrail" assertion was silently relying on Tremor Tabs rendering all panels at once. antd Tabs only renders the active tab's content, and defaultActiveKey is "submitted", so the button in the "Guardrails" tab wasn't in the DOM. Clicking the Guardrails tab first before asserting. --- ui/litellm-dashboard/src/components/guardrails.test.tsx | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/guardrails.test.tsx b/ui/litellm-dashboard/src/components/guardrails.test.tsx index 8cafc18eb9a..99c2474347e 100644 --- a/ui/litellm-dashboard/src/components/guardrails.test.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.test.tsx @@ -1,4 +1,4 @@ -import { render, screen } from "@testing-library/react"; +import { render, screen, fireEvent } from "@testing-library/react"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import GuardrailsPanel from "./guardrails"; import { getGuardrailsList } from "./networking"; @@ -40,6 +40,10 @@ vi.mock("./guardrails/GuardrailTestPlayground", () => ({ default: () =>
Mock Guardrail Test Playground
, })); +vi.mock("./guardrails/TeamGuardrailsTab", () => ({ + TeamGuardrailsTab: () =>
Mock Team Guardrails Tab
, +})); + vi.mock("@/utils/roles", () => ({ isAdminRole: vi.fn((role: string) => role === "admin"), })); @@ -99,6 +103,8 @@ describe("GuardrailsPanel", () => { it("should render the component", async () => { render(); expect(screen.getByText("Guardrails")).toBeInTheDocument(); + // Activate the Guardrails tab so its content (including the Add button) is rendered + fireEvent.click(screen.getByText("Guardrails")); expect(screen.getByText("+ Add New Guardrail")).toBeInTheDocument(); }); }); From e87f6cae5be7e3fb776d5062b5ae475dcf720975 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 4 Apr 2026 16:38:57 -0700 Subject: [PATCH 46/55] chore: poetry lock --- poetry.lock | 73 ++++++++++++++++++++++++++++++----------------------- 1 file changed, 41 insertions(+), 32 deletions(-) diff --git a/poetry.lock b/poetry.lock index dff1b8781d5..ca64f3101b1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.3.1 and should not be changed by hand. [[package]] name = "a2a-sdk" @@ -7,11 +7,11 @@ description = "A2A Python SDK" optional = false python-versions = ">=3.10" groups = ["main", "proxy-dev"] -markers = "python_version >= \"3.10\"" files = [ {file = "a2a_sdk-0.3.25-py3-none-any.whl", hash = "sha256:2fce38faea82eb0b6f9f9c2bcf761b0d78612c80ef0e599b50d566db1b2654b5"}, {file = "a2a_sdk-0.3.25.tar.gz", hash = "sha256:afda85bab8d6af0c5d15e82f326c94190f6be8a901ce562d045a338b7127242f"}, ] +markers = {main = "python_version >= \"3.10\" and extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] google-api-core = ">=1.26.0" @@ -386,6 +386,7 @@ files = [ {file = "azure_core-1.36.0-py3-none-any.whl", hash = "sha256:fee9923a3a753e94a259563429f3644aaf05c486d45b1215d098115102d91d3b"}, {file = "azure_core-1.36.0.tar.gz", hash = "sha256:22e5605e6d0bf1d229726af56d9e92bc37b6e726b141a18be0b4d424131741b7"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] requests = ">=2.21.0" @@ -406,6 +407,7 @@ files = [ {file = "azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c"}, {file = "azure_identity-1.25.3.tar.gz", hash = "sha256:ab23c0d63015f50b630ef6c6cf395e7262f439ce06e5d07a64e874c724f8d9e6"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] azure-core = ">=1.31.0" @@ -719,7 +721,7 @@ files = [ {file = "cffi-2.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:b882b3df248017dba09d6b16defe9b5c407fe32fc7c65a9c69798e6175601be9"}, {file = "cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529"}, ] -markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"proxy\"", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""} +markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"google\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""} [package.dependencies] pycparser = {version = "*", markers = "implementation_name != \"PyPy\""} @@ -1313,6 +1315,7 @@ files = [ {file = "cryptography-43.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff"}, {file = "cryptography-43.0.3.tar.gz", hash = "sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805"}, ] +markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"google\""} [package.dependencies] cffi = {version = ">=1.12", markers = "platform_python_implementation != \"PyPy\""} @@ -1378,7 +1381,7 @@ files = [ {file = "deprecated-1.3.1-py2.py3-none-any.whl", hash = "sha256:597bfef186b6f60181535a29fbe44865ce137a5079f295b479886c82729d5f3f"}, {file = "deprecated-1.3.1.tar.gz", hash = "sha256:b1b50e0ff0c1fddaa5708a2c6b0a6588bb09b892825ab2b214ac9ea9d92a5223"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] wrapt = ">=1.10,<3" @@ -2114,11 +2117,11 @@ description = "Google API client core library" optional = false python-versions = ">=3.7" groups = ["main", "proxy-dev"] -markers = "python_version >= \"3.14\"" files = [ {file = "google_api_core-2.25.2-py3-none-any.whl", hash = "sha256:e9a8f62d363dc8424a8497f4c2a47d6bcda6c16514c935629c257ab5d10210e7"}, {file = "google_api_core-2.25.2.tar.gz", hash = "sha256:1c63aa6af0d0d5e37966f157a77f9396d820fba59f9e43e9415bc3dc5baff300"}, ] +markers = {main = "python_version >= \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.14\""} [package.dependencies] google-auth = ">=2.14.1,<3.0.0" @@ -2146,7 +2149,7 @@ files = [ {file = "google_api_core-2.28.1-py3-none-any.whl", hash = "sha256:4021b0f8ceb77a6fb4de6fde4502cecab45062e66ff4f2895169e0b35bc9466c"}, {file = "google_api_core-2.28.1.tar.gz", hash = "sha256:2b405df02d68e68ce0fbc138559e6036559e685159d148ae5861013dc201baf8"}, ] -markers = {main = "(python_version >= \"3.10\" or extra == \"google\" or extra == \"extra-proxy\") and python_version <= \"3.13\"", proxy-dev = "python_version >= \"3.10\" and python_version <= \"3.13\""} +markers = {main = "python_version <= \"3.13\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.10\" and python_version <= \"3.13\""} [package.dependencies] google-auth = ">=2.14.1,<3.0.0" @@ -2183,7 +2186,7 @@ files = [ {file = "google_auth-2.49.1-py3-none-any.whl", hash = "sha256:195ebe3dca18eddd1b3db5edc5189b76c13e96f29e73043b923ebcf3f1a860f7"}, {file = "google_auth-2.49.1.tar.gz", hash = "sha256:16d40da1c3c5a0533f57d268fe72e0ebb0ae1cc3b567024122651c045d879b64"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] cryptography = ">=38.0.3" @@ -2351,11 +2354,11 @@ files = [ ] [package.dependencies] -google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]} -google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev" -grpc-google-iam-v1 = ">=0.12.4,<1.0.0dev" -proto-plus = ">=1.22.3,<2.0.0dev" -protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0dev" +google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0.dev0", extras = ["grpc"]} +google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0.dev0" +grpc-google-iam-v1 = ">=0.12.4,<1.0.0.dev0" +proto-plus = ">=1.22.3,<2.0.0.dev0" +protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0.dev0" [[package]] name = "google-cloud-resource-manager" @@ -2537,7 +2540,7 @@ files = [ {file = "googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038"}, {file = "googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\") or extra == \"google\" or extra == \"extra-proxy\""} [package.dependencies] grpcio = {version = ">=1.44.0,<2.0.0", optional = true, markers = "extra == \"grpc\""} @@ -2946,11 +2949,11 @@ description = "Consume Server-Sent Event (SSE) messages with HTTPX." optional = false python-versions = ">=3.9" groups = ["main", "proxy-dev"] -markers = "python_version >= \"3.10\"" files = [ {file = "httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc"}, {file = "httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d"}, ] +markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\")", proxy-dev = "python_version >= \"3.10\""} [[package]] name = "huey" @@ -3321,7 +3324,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.03.6" +jsonschema-specifications = ">=2023.3.6" referencing = ">=0.28.4" rpds-py = ">=0.7.1" @@ -3598,15 +3601,15 @@ files = [ [[package]] name = "litellm-proxy-extras" -version = "0.4.63" +version = "0.4.64" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_proxy_extras-0.4.63-py3-none-any.whl", hash = "sha256:46ec50083832b6b5ead86e53003657e1a53dc27bd95cbdbaee9c1343726e3acb"}, - {file = "litellm_proxy_extras-0.4.63.tar.gz", hash = "sha256:7161b27c3b38a840c13bb113b733196efafe2bb4d2ba22c9bd6e359c4b753aa2"}, + {file = "litellm_proxy_extras-0.4.64-py3-none-any.whl", hash = "sha256:e10f1d4bbfa84ce709e5ef559c8345d5bbe56e206655a5328384af87a078be6d"}, + {file = "litellm_proxy_extras-0.4.64.tar.gz", hash = "sha256:cca35fd41fea914dc067641df14937ba9037fa50177f11351f8c61902a434e92"}, ] [[package]] @@ -4095,6 +4098,7 @@ files = [ {file = "msal-1.35.1-py3-none-any.whl", hash = "sha256:8f4e82f34b10c19e326ec69f44dc6b30171f2f7098f3720ea8a9f0c11832caa3"}, {file = "msal-1.35.1.tar.gz", hash = "sha256:70cac18ab80a053bff86219ba64cfe3da1f307c74b009e2da57ef040eb1b5656"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] cryptography = ">=2.5,<49" @@ -4115,6 +4119,7 @@ files = [ {file = "msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca"}, {file = "msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] msal = ">=1.29,<2" @@ -4366,6 +4371,7 @@ files = [ {file = "nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9"}, {file = "nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f"}, ] +markers = {main = "extra == \"extra-proxy\""} [[package]] name = "numpy" @@ -4494,7 +4500,7 @@ files = [ {file = "opentelemetry_api-1.28.0-py3-none-any.whl", hash = "sha256:8457cd2c59ea1bd0988560f021656cecd254ad7ef6be4ba09dbefeca2409ce52"}, {file = "opentelemetry_api-1.28.0.tar.gz", hash = "sha256:578610bcb8aa5cdcb11169d136cc752958548fb6ccffb0969c1036b0ee9e5353"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] deprecated = ">=1.2.6" @@ -4600,7 +4606,7 @@ files = [ {file = "opentelemetry_sdk-1.28.0-py3-none-any.whl", hash = "sha256:4b37da81d7fad67f6683c4420288c97f4ed0d988845d5886435f428ec4b8429a"}, {file = "opentelemetry_sdk-1.28.0.tar.gz", hash = "sha256:41d5420b2e3fb7716ff4981b510d551eff1fc60eb5a95cf7335b31166812a893"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] opentelemetry-api = "1.28.0" @@ -4618,7 +4624,7 @@ files = [ {file = "opentelemetry_semantic_conventions-0.49b0-py3-none-any.whl", hash = "sha256:0458117f6ead0b12e3221813e3e511d85698c31901cac84682052adb9c17c7cd"}, {file = "opentelemetry_semantic_conventions-0.49b0.tar.gz", hash = "sha256:dbc7b28339e5390b6b28e022835f9bac4e134a80ebf640848306d3c5192557e8"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] deprecated = ">=1.2.6" @@ -5130,6 +5136,7 @@ files = [ {file = "prisma-0.11.0-py3-none-any.whl", hash = "sha256:22bb869e59a2968b99f3483bb417717273ffbc569fd1e9ceed95e5614cbaf53a"}, {file = "prisma-0.11.0.tar.gz", hash = "sha256:3f2f2fd2361e1ec5ff655f2a04c7860c2f2a5bc4c91f78ca9c5c6349735bf693"}, ] +markers = {main = "extra == \"extra-proxy\""} [package.dependencies] click = ">=7.1.2" @@ -5303,7 +5310,7 @@ files = [ {file = "proto_plus-1.26.1-py3-none-any.whl", hash = "sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66"}, {file = "proto_plus-1.26.1.tar.gz", hash = "sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version < \"3.13\" and (extra == \"extra-proxy\" or extra == \"google\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] protobuf = ">=3.19.0,<7.0.0" @@ -5331,7 +5338,7 @@ files = [ {file = "protobuf-5.29.5-py3-none-any.whl", hash = "sha256:6cf42630262c59b2d8de33954443d94b746c952b01434fc58a417fdbd2e84bd5"}, {file = "protobuf-5.29.5.tar.gz", hash = "sha256:bc1463bafd4b0929216c35f437a8e28731a2b7fe3d98bb77a600efced5a15c84"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\""} [[package]] name = "psutil" @@ -5491,7 +5498,7 @@ files = [ {file = "pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629"}, {file = "pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [[package]] name = "pyasn1-modules" @@ -5504,7 +5511,7 @@ files = [ {file = "pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a"}, {file = "pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] pyasn1 = ">=0.6.1,<0.7.0" @@ -5532,7 +5539,7 @@ files = [ {file = "pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934"}, {file = "pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2"}, ] -markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and implementation_name != \"PyPy\"", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""} +markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and implementation_name != \"PyPy\" and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"google\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""} [[package]] name = "pydantic" @@ -5755,6 +5762,7 @@ files = [ {file = "pyjwt-2.12.1-py3-none-any.whl", hash = "sha256:28ca37c070cad8ba8cd9790cd940535d40274d22f80ab87f3ac6a713e6e8454c"}, {file = "pyjwt-2.12.1.tar.gz", hash = "sha256:c74a7a2adf861c04d002db713dd85f84beb242228e671280bf709d765b03672b"}, ] +markers = {main = "extra == \"extra-proxy\" or extra == \"proxy\""} [package.dependencies] cryptography = {version = ">=3.4.0", optional = true, markers = "extra == \"crypto\""} @@ -6736,10 +6744,10 @@ files = [ ] [package.dependencies] -botocore = ">=1.37.4,<2.0a.0" +botocore = ">=1.37.4,<2.0a0" [package.extras] -crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"] +crt = ["botocore[crt] (>=1.37.4,<2.0a0)"] [[package]] name = "scikit-learn" @@ -6892,9 +6900,9 @@ tornado = ">=6.4.2,<7" urllib3 = ">=1.26,<3" [package.extras] -all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.00)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"] +all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.0)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"] bedrock = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)"] -cohere = ["cohere (>=5.9.4,<6.00)"] +cohere = ["cohere (>=5.9.4,<6.0)"] dev = ["dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "ipykernel (>=6.25.0,<7)", "mypy (>=1.7.1,<2)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"] docs = ["pydoc-markdown (>=4.8.2) ; python_version < \"3.12\""] fastembed = ["fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\""] @@ -7573,6 +7581,7 @@ files = [ {file = "tomlkit-0.13.3-py3-none-any.whl", hash = "sha256:c89c649d79ee40629a9fda55f8ace8c6a1b42deb912b2a8fd8d942ddadb606b0"}, {file = "tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1"}, ] +markers = {main = "extra == \"extra-proxy\""} [[package]] name = "tornado" @@ -8114,7 +8123,7 @@ files = [ {file = "wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22"}, {file = "wrapt-1.17.3.tar.gz", hash = "sha256:f66eb08feaa410fe4eebd17f2a2c8e2e46d3476e9f8c783daa8e09e0faa666d0"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [[package]] name = "wsproto" @@ -8309,4 +8318,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "9a2476d5f234f3ce45f399a77fb9e86bd0025e27e2bb905b0ecac7848c4a758c" +content-hash = "4964cafa67fee48aa1c7dd38ca08de20677b273d2b6faee2d6a264330f9099aa" From 61b295238b85a8ef1c7a1a3c5525ee31f7c3578a Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Sat, 4 Apr 2026 16:44:02 -0700 Subject: [PATCH 47/55] cherry-pick: tag query fix + MCP metadata support (#25145) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * added support for metadata (#24261) * added support for metadata * fix: PR review - meta truthiness, BlobResourceContents mimeType, add Blob+empty meta tests Made-with: Cursor * pyproject to .25 * feat(teams): resolve access group models/MCPs/agents in team endpoints Add access_group_models, access_group_mcp_server_ids, and access_group_agent_ids to /team/info and /v2/team/list responses. These fields contain resources inherited from access groups, kept separate from direct assignments so the UI can distinguish the source. Backend: _resolve_access_group_resources() helper resolves access group resources via existing _get_*_from_access_groups() functions. UI: Teams table and detail view show direct models as blue badges and access-group-sourced models as green badges. * perf(teams): single-pass access group resolution + asyncio.gather in list endpoint - Fetch each access group object once and extract all 3 resource fields in a single pass instead of 3 separate calls (3N → N lookups) - Use asyncio.gather to resolve access groups across teams concurrently in list_team_v2 instead of sequential awaits - Add 5 unit tests for _resolve_access_group_resources * docs: add default_team_params to config reference and update examples - Add default_team_params to litellm_settings reference table in config_settings.md with all sub-fields documented - Update self_serve.md and msft_sso.md examples to include team_member_permissions, tpm_limit, and rpm_limit - Fix misleading comment that implied default_team_params only applies to SSO auto-created teams — it applies to all /team/new calls * docs: clarify that models sub-field only applies to SSO auto-created teams * fix: lazy import get_access_object to break cyclic import + short-circuit all-proxy-models display - Remove get_access_object from module-level import in team_endpoints.py and use a lazy _get_access_object wrapper to avoid cyclic dependency - Add _prisma_client is None early-exit guard in _resolve_access_group_resources - Short-circuit UI to show "All Proxy Models" when team.models is empty or contains "all-proxy-models", skipping access group model resolution * add: making organizations a select instead of read only badges * fix(ui): only send organization_id when changed and use raw initial value * fix(ui): add paginated team search to usage page filter Replace the static team dropdown on the usage page with a new TeamMultiSelect component that uses the paginated v2/team/list endpoint with debounced server-side search and infinite scroll. * fix(ui): fix imports and update placeholder for team multi select * fix(ui): wire team_id filter to key alias dropdown on Virtual Keys tab The Key Alias dropdown on the Virtual Keys page was showing aliases from all teams regardless of which team was selected. The team_id was never passed through the frontend chain to the backend /key/aliases endpoint. - Backend: add optional team_id query param to /key/aliases endpoint - networking.tsx: add team_id param to keyAliasesCall - useKeyAliases: accept and forward team_id to API call and query key - filter.tsx: pass allFilters context to custom filter components - PaginatedKeyAliasSelect: read Team ID from allFilters and pass to hook * fix(tests): correct mock targets in TestResolveAccessGroupResources Three tests were patching the non-existent `get_access_object` instead of `_get_access_object` (the lazy-import wrapper), causing AttributeError. Also added missing `prisma_client` mock so tests get past the early-exit guard and actually exercise the resolution logic. * fix: use direct attribute access with or [] fallback in _resolve_access_group_resources Replace getattr(ag, "field", []) with ag.field or [] for cleaner access and safe handling if a field is None. * fix(ui): remove model source legend from team detail view The blue/green color distinction is self-explanatory; the legend added visual clutter without providing enough value. * fix(ui): add missing access_group fields to TeamData.team_info type The TeamData interface was missing access_group_models, access_group_mcp_server_ids, and access_group_agent_ids fields, causing a TypeScript build failure. * perf(teams): batch-fetch access groups in single DB query Replace per-ID _resolve_access_group_resources loop with a single find_many call that deduplicates IDs across all teams. Removes the N+1 query pattern on cold cache for the team list endpoint. * refactor(proxy): extract helpers to fix PLR0915 violations Extract `_apply_non_admin_alias_scope` from `key_aliases`, `_resolve_team_access_group_resources` from `team_info`, and `_enforce_list_team_v2_access` from `list_team_v2` to bring each function under ruff's 50-statement limit. No behavior changes. * test(ui): update tests to match new team_id / access-group signatures - useKeyAliases, PaginatedKeyAliasSelect: add trailing `undefined` to spy matchers for the new `team_id` param on `useInfiniteKeyAliases` and `keyAliasesCall`. - EntityUsage: mock new `TeamMultiSelect` child so QueryClientProvider is not required for team-entity tests. - ModelsCell: replace the overflow-accordion test with one that verifies the new collapse-on-`all-proxy-models` behavior (no accordion, single badge). * fix(ui): send null (not '') for cleared organization_id on team update AntD