diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 12c3c5e2be9..b288269b0a2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6164,7 +6164,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4": { "cache_read_input_token_cost": 2.8e-07, @@ -6199,7 +6202,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4": { "cache_read_input_token_cost": 2.8e-07, @@ -6234,7 +6240,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -6276,7 +6285,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, @@ -6312,7 +6324,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, @@ -6348,7 +6363,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -7301,8 +7319,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, @@ -7337,8 +7355,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano": { "cache_read_input_token_cost": 2e-08, @@ -7372,8 +7390,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano-2026-03-17": { "cache_read_input_token_cost": 2e-08, @@ -7408,8 +7426,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-image-1": { "cache_read_input_token_cost": 1.25e-06, @@ -8712,6 +8730,268 @@ "/v1/images/generations" ] }, + "azure_ai/FW-DeepSeek-V3.2": { + "cache_read_input_token_cost": 3.1e-07, + "input_cost_per_token": 6.2e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1.85e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-DeepSeek-V4-Pro": { + "cache_read_input_token_cost": 1.65e-07, + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.828e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5": { + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 1.1e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.52e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.1": { + "cache_read_input_token_cost": 2.86e-07, + "input_cost_per_token": 1.54e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 202800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.2": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.54e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.2-Fast": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-Inkling": { + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://fireworks.ai/models/fireworks/inkling", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-Kimi-K2.5": { + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6.6e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.3e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Kimi-K2.6": { + "cache_read_input_token_cost": 1.76e-07, + "input_cost_per_token": 1.045e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Kimi-K2.7-Code": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 1.05e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Kimi-K3": { + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-kimi-k3-through-fireworks-ai-on-microsoft-foundry/4540187", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-MiniMax-M2.5": { + "cache_read_input_token_cost": 3.3e-08, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-MiniMax-M3": { + "cache_read_input_token_cost": 6.6e-08, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "max_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Nemotron-3-Ultra-NVFP4": { + "cache_read_input_token_cost": 1.19e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://fireworks.ai/models/fireworks/nemotron-3-ultra-nvfp4", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/MAI-Image-2.5": { "input_cost_per_image_token": 8e-06, "input_cost_per_token": 5e-06, @@ -9329,6 +9609,24 @@ "supports_tool_choice": true, "supports_web_search": true }, + "azure_ai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-grok-4-3-on-microsoft-foundry-latest-generation-agentic-capabilities/4517096", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, @@ -24869,7 +25167,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -27711,6 +28012,93 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.2": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.2-contributor": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1e-07, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index e56e74dba8c..9103f9fbbf5 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -554,13 +554,20 @@ async def get_all_mcp_servers( ) -> list[LiteLLM_MCPServerTable]: """ Returns mcp servers from the db, optionally filtered by approval_status. - Pass approval_status=None to return all servers regardless of approval state. + Pass approval_status=None to return every server except drafts, which back the admin OAuth + session flow, are addressable only by their own server_id, and must never appear in a listing. + NULL approval_status predates the approval workflow, so those rows are kept explicitly rather + than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them. """ try: - where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = {} - if approval_status is not None: - where["approval_status"] = approval_status - mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where if where else {}) + where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = ( + {"approval_status": approval_status} + if approval_status is not None + # mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop + # NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts + else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]} + ) + mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where) tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] for table in tables: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b1e444e55d6..0239368ad0e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -65,6 +65,7 @@ from litellm.proxy.auth.budget_throttle import ( should_throttle_budget_exceeded, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, @@ -2016,6 +2017,44 @@ async def _cache_team_object( ) +async def delete_cache_team_object( + team_id: str, + team_alias: str | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> None: + """ + Evict both keys `_cache_team_object` writes. + + `get_team_object` reads the id key and the JWT `team_alias_jwt_field` path reads the alias key, + so leaving either behind keeps a deleted team resolvable for auth until its TTL expires. + + Mirrors `delete_cached_project_object`: evicting locally only reaches the worker handling the + delete, so every key is also broadcast to drop the other workers' in-memory copies. + + Eviction is best-effort, matching `_cache_team_object`. `delete_team` calls this after the team + rows are already gone, so letting an unreachable cache backend raise here would fail a request + whose delete has committed. + """ + keys: Final = (f"team_id:{team_id}", *((f"team_alias:{team_alias}",) if team_alias else ())) + + for key in keys: + try: + user_api_key_cache.delete_cache(key=key) + + ## UPDATE REDIS CACHE ## + if proxy_logging_obj is not None: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete + verbose_proxy_logger.warning( + "Failed to invalidate cached team entry %s on delete; " + "a deleted team may be served until its TTL expires: %s", + key, + e, + ) + await publish_auth_cache_invalidation(cache_key=key) + + async def _cache_key_object( hashed_token: str, user_api_key_obj: UserAPIKeyAuth, @@ -2051,6 +2090,61 @@ async def _delete_cache_key_object( await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) +class TeamNotFoundError(HTTPException): + """The team row is provably absent, as opposed to merely unreadable. + + ``get_team_object`` reports every failure as a 404, so a deleted team and a + database that would not answer are indistinguishable to its callers. Callers + that must not treat a degraded read as a definitive answer, such as the + authorization fallback in ``user_api_key_auth``, key on this subclass. It + stays a 404 carrying the same detail, so every other caller is unaffected. + """ + + def __init__(self, team_id: str) -> None: + super().__init__( + status_code=404, + detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."}, + ) + + +async def delete_cache_key_objects( + hashed_tokens: Sequence[str], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> None: + """ + Evict a batch of key objects, for callers that delete keys in bulk rather than through + `/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left + cached after its row is gone keeps buying access until its TTL expires. + + Evicting locally only reaches this worker, so each token is also broadcast: a deleted key left + in a peer worker's in-memory cache still authenticates there until its TTL expires. + + Best-effort per key: the rows are already deleted by the time this runs, so an unreachable + cache backend must not abort the caller partway through its own cascade. + """ + results: Final = await asyncio.gather( + *( + _delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for hashed_token in hashed_tokens + ), + return_exceptions=True, + ) + + for hashed_token, result in zip(hashed_tokens, results): + if isinstance(result, BaseException): + verbose_proxy_logger.warning( + "Failed to evict cached key entry for %s; a deleted key may authenticate until its TTL expires: %s", + hashed_token, + result, + ) + await publish_auth_cache_invalidation(cache_key=hashed_token) + + @log_db_metrics async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None @@ -2096,6 +2190,10 @@ async def _get_team_object_from_user_api_key_cache( ) if should_check_db: response = await _get_team_db_check(team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert) + # The database answered and the row is not there. Distinct from every + # other failure here, which leaves the team's grant unknown. + if response is None: + raise TeamNotFoundError(team_id=team_id) else: response = None @@ -2217,6 +2315,8 @@ async def get_team_object( key=key, team_id_upsert=team_id_upsert, ) + except TeamNotFoundError: + raise except Exception: raise HTTPException( status_code=404, @@ -2556,6 +2656,8 @@ class ExperimentalUIJWTToken: user_info: LiteLLM_UserTable, team_id: str | None = None, team_alias: str | None = None, + team_models: Sequence[str] | None = None, + team_model_aliases: Mapping[str, str] | None = None, max_budget: float | None = None, ) -> str: """ @@ -2568,6 +2670,8 @@ class ExperimentalUIJWTToken: user_info: User information from the database team_id: Team ID for the user (optional, uses user's team if available) team_alias: Team alias for the selected team, if available + team_models: Model allowlist granted by the selected team + team_model_aliases: Team model aliases for the selected team Returns: Encrypted JWT token string @@ -2606,7 +2710,9 @@ class ExperimentalUIJWTToken: user_id=user_info.user_id, team_id=_team_id, team_alias=team_alias, - models=user_info.models, + team_models=list(team_models) if team_models is not None else [], + team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None, + models=[] if _team_id is not None else user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), is_session_token=True, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f7a04ba79e7..39e1c14a6e6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, + TeamNotFoundError, _cache_key_object, _can_object_call_model, _check_end_user_budget, @@ -85,6 +86,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( PrismaClient, @@ -2161,6 +2163,28 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) +def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseException) -> bool: + """Whether the token's own team fields may stand in for a team that failed to + resolve, without widening access. + + A team that is provably gone is a definitive answer, not a degraded read, so + nothing may stand in for it and no setting may override that. + + Otherwise the team's grant is merely unknown. A token carrying one may vouch, + since replaying a recorded grant cannot widen it and denying every team key + while the row is briefly unreadable would trade the widening for an outage. A + token carrying none may not: ``team_models=[]`` reads as every model and + ``team_blocked=False`` as unblocked. ``allow_requests_on_db_unavailable`` opts + back out, and is only consulted here because the failure is known by this + point to be a degraded read. + """ + if isinstance(lookup_error, TeamNotFoundError): + return False + if valid_token.team_models: + return True + return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + + @tracer.wrap() async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, @@ -2364,7 +2388,12 @@ async def _run_centralized_common_checks( if isinstance(team_result, BaseException): # Token-derived fallback only valid when a team_id is set; # _team_obj_from_token asserts that precondition. - team_object = _team_obj_from_token(user_api_key_auth_obj) if user_api_key_auth_obj.team_id is not None else None + if user_api_key_auth_obj.team_id is None: + team_object = None + elif _token_can_vouch_for_team(user_api_key_auth_obj, team_result): + team_object = _team_obj_from_token(user_api_key_auth_obj) + else: + raise team_result else: team_object = team_result diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index de9d38963c1..72ed67728d8 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -36,6 +36,17 @@ The base URL is resolved in this order of precedence: 3. `base_url` from `~/.litellm/config.json` 4. `http://localhost:4000` +### Hiding commands from the listings + +Deployments that hand `lite` to end users often want to advertise only part of it. Store the commands to keep out of the listings, comma separated: + +```bash +lite config set hidden_commands codex,opencode +lite config unset hidden_commands # list everything again +``` + +Hidden commands drop out of both `lite --help` and the interactive shell's "Available commands" block, and stay runnable so existing scripts keep working + ## Global Options - `--version`, `-v`: Print the LiteLLM Proxy client and server version and exit. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index dfc70a8df7c..ed2bf2be03d 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -1,5 +1,6 @@ import os import shutil +import subprocess import sys from collections.abc import Callable, Mapping, Sequence from typing import Final @@ -142,8 +143,95 @@ def verify_proxy_key( ) -def _exec(path: str, args: Sequence[str], env: Mapping[str, str]) -> None: - os.execvpe(path, list(args), dict(env)) +_WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"}) +_CMD_PERCENT_GUARD: Final = "%%cd:~,%" +_CMD_LINE_BREAKS: Final = ("\r", "\n") + + +def _double_trailing_backslashes(segment: str) -> str: + bare: Final = segment.rstrip("\\") + return bare + "\\" * 2 * (len(segment) - len(bare)) + + +def _quote_for_cmd(token: str) -> str: + """Quote one token so both parsers that read it see the original text. + + Follows the algorithm the Rust standard library settled on for batch files + after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a + quoted string on a lone `"` and so wants an embedded one doubled, and the + shim's own interpreter, which re-splits `%*` under C runtime rules where a + backslash escapes the quote that follows it, so every backslash run standing + before a quote is doubled. Quoting cannot stop cmd expanding `%VAR%`, so each + `%` is prefixed with `%%cd:~,`: the zero-length substring of the always + defined `cd` expands to nothing and leaves no `%` pair for cmd to match. + """ + escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"')) + return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"' + + +def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: + """Build what CreateProcess runs, routing batch shims through cmd.exe. + + npm installs Claude Code as `claude.cmd`, which PATHEXT lets shutil.which + resolve but CreateProcess refuses to run (WinError 193), so a shim has to go + through the command processor. cmd.exe does not follow the C runtime quoting + that subprocess would apply to an argument list, and it would split on `&` or + `|` in a forwarded argument, so the shim case is emitted as one verbatim + command line with every token quoted. Every switch is load-bearing: `/s` + makes cmd strip only the outer pair, leaving each token quoted and its + metacharacters inert, `/e:on` keeps the command extensions that the percent + guard is built out of, `/v:off` keeps `!` from expanding, and `/d` keeps a + machine's AutoRun commands out of the launch. argv[0] carries the + caller-facing name on POSIX; Windows needs the resolved path there. + + Raises AgentRunError for an argument holding a line break, which cmd would + read as the end of the command line and silently drop the rest of. + """ + rest: Final = tuple(args[1:]) + if os.path.splitext(path)[1].lower() not in _WINDOWS_SHIM_SUFFIXES: + return (path, *rest) + if any(brk in token for token in rest for brk in _CMD_LINE_BREAKS): + raise AgentRunError( + f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on " + "Windows: cmd.exe ends the command line there, so the agent would silently lose it." + ) + inner: Final = " ".join(_quote_for_cmd(token) for token in (path, *rest)) + return f'cmd.exe /d /e:on /v:off /s /c "{inner}"' + + +def _spawn_and_wait(command: str | Sequence[str], env: Mapping[str, str]) -> int: + return subprocess.run(command, env=dict(env), check=False).returncode + + +def _replace_process( + path: str, + args: Sequence[str], + env: Mapping[str, str], + *, + execvpe: Callable[..., None] = os.execvpe, +) -> None: + execvpe(path, list(args), dict(env)) + + +def _hand_off( + path: str, + args: Sequence[str], + env: Mapping[str, str], + *, + platform: str = sys.platform, + replace: Callable[[str, Sequence[str], Mapping[str, str]], None] = _replace_process, + spawn: Callable[[str | Sequence[str], Mapping[str, str]], int] = _spawn_and_wait, +) -> None: + """Replace this process with the agent; on Windows, run it as a child instead. + + os.exec* has no process-replacement semantics on Windows: the C runtime + spawns a detached child and terminates the parent, so the shell reclaims the + console and the agent's TUI never gets one. Windows therefore waits on the + child and exits with its status. + """ + if platform.startswith("win"): + raise SystemExit(spawn(_windows_command(path, args), env)) + replace(path, list(args), dict(env)) def _restore_controlling_terminal() -> None: @@ -175,13 +263,14 @@ def run_agent( base_env: Mapping[str, str] | None = None, which: Callable[[str], str | None] = shutil.which, verify: Callable[[str, str], None] = verify_proxy_key, - launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec, + launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, reattach_terminal: Callable[[], None] | None = None, ) -> None: """Validate, wire the environment, and hand off to the agent. - On success this replaces the current process and never returns. Raises - AgentRunError for missing binaries, an unreachable proxy, or a rejected key. + On success this never returns: POSIX replaces the current process, Windows + waits on the agent and exits with its status. Raises AgentRunError for + missing binaries, an unreachable proxy, or a rejected key. reattach_terminal, when given, runs just before handoff to restore stdin. """ if not command: @@ -277,9 +366,9 @@ def _make_agent_command(binary: str, display_name: str) -> click.Command: return _command -def agent_commands() -> list[click.Command]: +def agent_commands() -> tuple[click.Command, ...]: """Build one top-level command per known agent, e.g. `lite claude`.""" - return [_make_agent_command(binary, name) for binary, (name, _profiles) in _KNOWN_AGENTS.items()] + return tuple(_make_agent_command(binary, name) for binary, (name, _profiles) in _KNOWN_AGENTS.items()) __all__ = [ diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 8f1fcac740a..19dd407ba19 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -1,8 +1,9 @@ import json import os import sys -from collections.abc import Mapping +from collections.abc import Callable, Mapping from pathlib import Path +from types import MappingProxyType from typing import Final from urllib.parse import urlparse @@ -11,7 +12,7 @@ from pydantic import TypeAdapter from .private_json import write_private_json -ALLOWED_CONFIG_KEYS: Final[tuple[str, ...]] = ("base_url",) +HIDDEN_COMMANDS_KEY: Final = "hidden_commands" _config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str]) @@ -49,6 +50,48 @@ def get_config_value(key: str) -> str | None: return load_config().get(key) +def parse_hidden_commands(raw: str | None) -> frozenset[str]: + """Split a stored `hidden_commands` value, e.g. "codex, opencode".""" + return frozenset(name.strip() for name in (raw or "").split(",") if name.strip()) + + +def hidden_command_names() -> frozenset[str]: + """Top-level commands the operator chose to keep out of `lite`'s listings.""" + return parse_hidden_commands(get_config_value(HIDDEN_COMMANDS_KEY)) + + +def _normalize_base_url(value: str) -> str: + parsed: Final = urlparse(value) + if parsed.scheme not in ("http", "https") or not parsed.netloc: + raise click.UsageError("base_url must be a full http:// or https:// URL including a host") + if "?" in value or "#" in value: + raise click.UsageError("base_url must not include a query string or fragment") + return value.rstrip("/") + + +def _normalize_hidden_commands(value: str) -> str: + names: Final = parse_hidden_commands(value) + if not names: + raise click.UsageError( + f"{HIDDEN_COMMANDS_KEY} must be a comma-separated list of command names, e.g. " + f"`lite config set {HIDDEN_COMMANDS_KEY} codex,opencode`. To list everything again, " + f"run `lite config unset {HIDDEN_COMMANDS_KEY}`" + ) + if any(" " in name for name in names): + raise click.UsageError(f"{HIDDEN_COMMANDS_KEY} entries must be single command names, without spaces") + return ",".join(sorted(names)) + + +_NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType( + { + "base_url": _normalize_base_url, + HIDDEN_COMMANDS_KEY: _normalize_hidden_commands, + } +) + +ALLOWED_CONFIG_KEYS: Final[tuple[str, ...]] = tuple(_NORMALIZERS) + + @click.group(name="config") def config_commands() -> None: """Manage persistent CLI configuration (~/.litellm/config.json)""" @@ -59,17 +102,11 @@ def config_commands() -> None: @click.argument("value") def set_config(key: str, value: str) -> None: """Set a config KEY to VALUE (e.g. `lite config set base_url https://your-proxy.example.com`)""" - if key not in ALLOWED_CONFIG_KEYS: + normalizer: Final = _NORMALIZERS.get(key) + if normalizer is None: raise click.UsageError(f"Unknown config key '{key}'. Allowed keys: {', '.join(ALLOWED_CONFIG_KEYS)}") - if key == "base_url": - parsed: Final = urlparse(value) - if parsed.scheme not in ("http", "https") or not parsed.netloc: - raise click.UsageError("base_url must be a full http:// or https:// URL including a host") - if "?" in value or "#" in value: - raise click.UsageError("base_url must not include a query string or fragment") - - normalized_value: Final = value.rstrip("/") + normalized_value: Final = normalizer(value) save_config({**load_config(), key: normalized_value}) click.echo(f"Set {key} = {normalized_value} in {get_config_file_path()}") diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index 84862bb4a55..b4f44240adb 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -74,8 +74,9 @@ def styled_prompt(): def show_commands(): - """Display available commands.""" + """Display available commands, minus any the operator chose to hide.""" from .commands.agents import agent_commands + from .commands.config import hidden_command_names commands = [ ("login", "Authenticate with the LiteLLM proxy server"), @@ -96,9 +97,12 @@ def show_commands(): ("quit", "Exit the interactive session"), ] + hidden: Final = hidden_command_names() + click.echo("Available commands:") for cmd, description in commands: - click.echo(f" {cmd:<20} {description}") + if cmd not in hidden: + click.echo(f" {cmd:<20} {description}") click.echo() diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 95d4a751226..3a289736c66 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -12,7 +12,7 @@ from .commands.agents import agent_commands from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami from .commands.autoroute.commands import autoroute_group from .commands.chat import chat -from .commands.config import config_commands, get_config_value +from .commands.config import config_commands, get_config_value, hidden_command_names from .commands.credentials import credentials from .commands.encryption import encryption from .commands.http import http @@ -43,7 +43,21 @@ def print_version(base_url: str, api_key: str | None): click.echo(f"Could not retrieve server version: {e}") -@click.group(invoke_without_command=True) +class HideConfiguredCommandsGroup(click.Group): + """Group that omits operator-hidden commands from listings, still running them. + + Deployments hand `lite` to users who should only see a curated subset of + commands (`lite config set hidden_commands codex,opencode`). Filtering the + listing rather than dropping the commands keeps anyone's existing scripts + working. + """ + + def list_commands(self, ctx: click.Context) -> list[str]: + hidden: Final = hidden_command_names() + return [name for name in super().list_commands(ctx) if name not in hidden] + + +@click.group(cls=HideConfiguredCommandsGroup, invoke_without_command=True) @click.option( "--version", "-v", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a416a197ab8..6e1e6d22cb1 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2311,7 +2311,7 @@ async def delete_user( fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}}) teams_to_update = [] for team in fetch_all_teams: - is_member_in_team, new_team_members = _cleanup_members_with_roles( + removed_team_members, new_team_members = _cleanup_members_with_roles( existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()), data=TeamMemberDeleteRequest( team_id=team.team_id, @@ -2319,7 +2319,7 @@ async def delete_user( user_email=user_row.user_email, ), ) - if is_member_in_team: + if removed_team_members: _db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members] team.members_with_roles = json.dumps(_db_new_team_members) teams_to_update.append(team) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae1bb52278a..b72b5d218b2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -77,6 +77,8 @@ from litellm.proxy.auth.auth_checks import ( _cache_team_object, allowed_route_check_inside_route, can_org_access_model, + delete_cache_key_objects, + delete_cache_team_object, get_org_object, get_team_membership, get_team_object, @@ -313,6 +315,11 @@ class _TeamIdInFilter(TypedDict, total=False): team_id: Mapping[str, Sequence[str]] +_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """ +UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams) +""" + + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) @@ -2657,6 +2664,11 @@ async def _add_team_members_to_team( serialize on the row lock and each appends onto the other's committed result, instead of both rewriting the whole JSON array from a stale snapshot (which silently drops one member on the losing write). + + The same lock serializes this against /team/delete: the delete cannot remove + the row while the reconcile holds it, and a reconcile that finds the row + already gone cleans up after itself rather than leaving the member pointing + at a deleted team id. """ # Process and add new members updated_users, updated_team_memberships = await _process_team_members( @@ -2667,11 +2679,42 @@ async def _add_team_members_to_team( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - async with prisma_client.tx() as tx: - complete_team_data.members_with_roles = await TeamRepository(prisma_client).get_members_with_roles_locked( - tx, data.team_id + updated_team: Final = await _write_members_with_roles_locked( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + updated_users=updated_users, + ) + if updated_team is None: + await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client) + raise HTTPException( + status_code=404, + detail={"error": f"Team={data.team_id} was deleted while this member add was running"}, ) + return updated_team, updated_users, updated_team_memberships + + +async def _write_members_with_roles_locked( + data: TeamMemberAddRequest, + complete_team_data: LiteLLM_TeamTable, + prisma_client: PrismaClient, + updated_users: list[LiteLLM_UserTable], +) -> LiteLLM_TeamTable | None: + """Reconcile members_with_roles under the team row lock. None when the team row is gone. + + That read is at least as recent as the user and membership writes the caller + already made, so a missing row means /team/delete committed after them. Its + post-delete sweep can have run before those writes landed, which is why the + caller sweeps this team id again rather than only reporting the 404. + """ + async with prisma_client.tx() as tx: + locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id) + if locked_members is None: + return None + + complete_team_data.members_with_roles = locked_members + await _update_team_members_list( data=data, complete_team_data=complete_team_data, @@ -2679,13 +2722,11 @@ async def _add_team_members_to_team( ) _db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team: Final = await tx.litellm_teamtable.update( + return await tx.litellm_teamtable.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, ) - return updated_team, updated_users, updated_team_memberships - def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: """Update the Prometheus team members gauge after a membership change. @@ -3091,26 +3132,27 @@ async def team_member_add( ) +def _is_member_addressed_by(member: Member, data: TeamMemberDeleteRequest) -> bool: + return (data.user_id is not None and member.user_id is not None and data.user_id == member.user_id) or ( + data.user_email is not None and member.user_email is not None and data.user_email == member.user_email + ) + + def _cleanup_members_with_roles( existing_team_row: LiteLLM_TeamTable, data: TeamMemberDeleteRequest, -) -> tuple[bool, list[Member]]: - """Cleanup members_with_roles list for a team.""" - is_member_in_team = False - new_team_members: Final[list[Member]] = [] - for m in existing_team_row.members_with_roles: - if ( - data.user_id is not None - and m.user_id is not None - and data.user_id == m.user_id - or data.user_email is not None - and m.user_email is not None - and data.user_email == m.user_email - ): - is_member_in_team = True - continue - new_team_members.append(m) - return is_member_in_team, new_team_members +) -> tuple[tuple[Member, ...], list[Member]]: + """Split a team's members_with_roles into the entries the request addresses and the ones that stay. + + The addressed entries are returned rather than a bare found/not-found flag because they carry the + user_id the request may not have supplied, and every cleanup that keys off the user rather than + off the roster has to run against that id. + """ + removed_team_members: Final = tuple( + m for m in existing_team_row.members_with_roles if _is_member_addressed_by(m, data) + ) + new_team_members: Final = [m for m in existing_team_row.members_with_roles if not _is_member_addressed_by(m, data)] + return removed_team_members, new_team_members @router.post( @@ -3182,12 +3224,12 @@ async def team_member_delete( ) ## DELETE MEMBER FROM TEAM - is_member_in_team, new_team_members = _cleanup_members_with_roles( + removed_team_members, new_team_members = _cleanup_members_with_roles( existing_team_row=existing_team_row, data=data, ) - if not is_member_in_team: + if not removed_team_members: raise HTTPException(status_code=400, detail={"error": "User not found in team"}) existing_team_row.members_with_roles = new_team_members @@ -3205,38 +3247,28 @@ async def team_member_delete( ## DELETE TEAM ID from USER ROW, IF EXISTS ## # get user row - key_val: Final = {} - if data.user_id is not None: - key_val["user_id"] = data.user_id - elif data.user_email is not None: - key_val["user_email"] = data.user_email - existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many( - where=key_val + removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None) + key_val: Final[Mapping[str, object]] = ( + {"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email} ) + existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val) - if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): - for existing_user in existing_user_rows: - team_list = [] - if data.team_id in existing_user.teams: - team_list = existing_user.teams - team_list.remove(data.team_id) - await _user_db(prisma_client).update( - where={ - "user_id": existing_user.user_id, - }, - data={"teams": {"set": team_list}}, - ) + for existing_user in existing_user_rows: + if data.team_id in existing_user.teams: + await _user_db(prisma_client).update( + where={ + "user_id": existing_user.user_id, + }, + data={"teams": {"set": [team for team in existing_user.teams if team != data.team_id]}}, + ) # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = set[str]() - if data.user_id is not None: - user_ids_to_delete.add(data.user_id) - if existing_user_rows is not None and isinstance(existing_user_rows, list): - for existing_user in existing_user_rows: - if getattr(existing_user, "user_id", None): - user_ids_to_delete.add(existing_user.user_id) + user_ids_to_delete: Final = removed_user_ids.union( + (data.user_id,) if data.user_id is not None else (), + (user.user_id for user in existing_user_rows if user.user_id), + ) - for _uid in user_ids_to_delete: + for _uid in sorted(user_ids_to_delete): await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM @@ -3248,7 +3280,7 @@ async def team_member_delete( # Fetch keys before deletion to persist them keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={ - "user_id": {"in": list(user_ids_to_delete)}, + "user_id": {"in": sorted(user_ids_to_delete)}, "team_id": data.team_id, } ) @@ -3263,7 +3295,7 @@ async def team_member_delete( await _tokens_db(prisma_client).delete_many( where={ - "user_id": {"in": list(user_ids_to_delete)}, + "user_id": {"in": sorted(user_ids_to_delete)}, "team_id": data.team_id, } ) @@ -3662,6 +3694,8 @@ async def delete_team( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, + proxy_logging_obj, + user_api_key_cache, ) if prisma_client is None: @@ -3755,6 +3789,12 @@ async def delete_team( await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") + await _invalidate_deleted_key_cache( + keys=keys_to_delete, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ## DELETE ASSOCIATED BYOK MODELS # Runs before the team rows are deleted so a mid-flight failure never leaves # the team gone with its models orphaned. @@ -3788,11 +3828,90 @@ async def delete_team( ) await asyncio.gather(*tasks) + await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + ## DELETE TEAMS deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team") + + # Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and + # `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a + # concurrent auth lookup re-caches the still-present team and the delete looks like it never + # invalidated anything. Nothing fallible runs between the delete and this, or a failure there + # would strand the deleted team in cache. + await _invalidate_deleted_team_cache( + teams=team_rows, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep + # and the delete would have re-appended the reference; an add still in flight sees the row + # missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and + # keeping the first one means a failure here still leaves a team the admin can retry deleting. + await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + return deleted_teams +async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None: + """ + Strip the deleted team ids from every user row and team-membership row that still references them. + + The per-member `team_member_delete` pass above only reaches users listed in the team's + `members_with_roles`, so a user row that outlived its roster entry is invisible to it and keeps + surfacing the team on `/user/info` after the team is gone. + + #36839 closed the route that created that drift, by resolving member removal off the roster + entry's `user_id` rather than the identifier the caller happened to pass. It does not backfill + rows that already drifted, which is the state this was reported against, so the sweep still has + to run on delete. + + `array_remove` rather than read-filter-write: rewriting the whole array from a snapshot read + outside a transaction drops any team a concurrent `/team/member_add` appended in between. + """ + for team_id in team_ids: + _ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id) + + _ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)})) + + +async def _invalidate_deleted_key_cache( + keys: Sequence[LiteLLM_VerificationToken], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Evict the auth cache entry for every key deleted along with the team. + + `/key/delete` evicts as it goes, but the bulk delete above writes straight to the db. Auth + resolves a cached key object without re-reading the team, so a key belonging to a deleted team + keeps buying access until its TTL expires. + """ + await delete_cache_key_objects( + hashed_tokens=tuple(key.token for key in keys), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _invalidate_deleted_team_cache( + teams: Sequence[LiteLLM_TeamTable], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + _ = await asyncio.gather( + *( + delete_cache_team_object( + team_id=team.team_id, + team_alias=team.team_alias, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for team in teams + ) + ) + + def _transform_teams_to_deleted_records( teams: list[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a2c50590dd5..b87ad8597dc 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -21,6 +21,7 @@ from copy import deepcopy from html import escape from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -40,6 +41,7 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -185,6 +187,7 @@ class _PrismaTableActions(Protocol[_DbRecordT]): async def find_many( self, where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, ) -> Sequence[_DbRecordT]: ... async def update( @@ -241,6 +244,45 @@ def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDe return repo.table +_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) + + +def _decode_model_aliases(value: object) -> object: + """``/team/new`` stores team model aliases as a JSON-encoded string in the Json column.""" + if not isinstance(value, str): + return value + try: + return _MODEL_ALIASES_ADAPTER.validate_json(value) + except ValidationError: + return None + + +class _TeamModelAliasTable(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + + model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None + + +class _TeamRowGrants(BaseModel): + team_id: str + team_alias: str | None = None + models: tuple[str, ...] = () + litellm_model_table: _TeamModelAliasTable | None = None + + +class _CliSsoTeamDetail(BaseModel): + """The per-team snapshot cached in the CLI SSO flow and echoed to the CLI on poll.""" + + team_id: str | None = None + team_alias: str | None = None + team_models: tuple[str, ...] + team_model_aliases: Mapping[str, str] | None = None + + +_CLI_SSO_TEAM_DETAILS_ADAPTER: Final = TypeAdapter(tuple[_CliSsoTeamDetail, ...]) +_TEAMLESS_CLI_SSO_TEAM_DETAIL: Final = _CliSsoTeamDetail(team_models=()) + + class _CustomSsoCall(Protocol): async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... @@ -2147,27 +2189,55 @@ async def _build_cli_sso_user_defined_values( ) +def _cli_sso_team_detail(team_row: Mapping[str, object]) -> _CliSsoTeamDetail: + team: Final = _TeamRowGrants.model_validate(team_row) + alias_table: Final = team.litellm_model_table + return _CliSsoTeamDetail( + team_id=team.team_id, + team_alias=team.team_alias, + team_models=team.models, + team_model_aliases=alias_table.model_aliases if alias_table is not None else None, + ) + + async def _fetch_cli_sso_team_details( prisma_client: PrismaClient, teams: Sequence[str], -) -> list[dict[str, object]]: - team_details: Final[list[dict[str, object]]] = [] +) -> tuple[_CliSsoTeamDetail, ...] | None: + """``None`` means the lookup itself failed, which is not the same as the user having no teams.""" + if not teams: + return () try: - if teams: - prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( - where={"team_id": {"in": teams}} - ) - for team_row in prisma_teams: - team_dict = team_row.model_dump() - team_details.append( - { - "team_id": team_dict.get("team_id"), - "team_alias": team_dict.get("team_alias"), - } - ) + prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( + where={"team_id": {"in": teams}}, + include={"litellm_model_table": True}, + ) except Exception as e: verbose_proxy_logger.error("Error fetching team details for CLI SSO session: %s", e) - return team_details + return None + return tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams) + + +def _cli_sso_session_teams(team_details: Sequence[_CliSsoTeamDetail]) -> list[str]: + """The teams a login may bind to: only those whose row still exists. + + A team deleted out from under a membership, which is what deleting an organization + leaves behind, can never resolve its grants, so offering it would refuse every + future login for that user with nothing they could do to recover. + """ + return [detail.team_id for detail in team_details if detail.team_id is not None] + + +def _selected_cli_sso_team_detail(team_details: object, team_id: str | None) -> _CliSsoTeamDetail | None: + """``None`` means the team's grants are unknown. An empty grant is a real value meaning unrestricted, + so an unknown one must not be minted as empty.""" + if team_id is None: + return _TEAMLESS_CLI_SSO_TEAM_DETAIL + try: + details: Final = _CLI_SSO_TEAM_DETAILS_ADAPTER.validate_python(team_details) + except ValidationError: + return None + return next((detail for detail in details if detail.team_id == team_id), None) async def _complete_cli_sso_callback_session( @@ -2210,6 +2280,12 @@ async def _complete_cli_sso_callback_session( teams = user_info.teams if isinstance(user_info.teams, list) else [] team_details: Final = await _fetch_cli_sso_team_details(prisma_client=prisma_client, teams=teams) + if team_details is None: + raise HTTPException( + status_code=500, + detail="Could not resolve team model grants for this login. Please try again", + ) + resolved_teams: Final = _cli_sso_session_teams(team_details) attribution_metadata: Final = build_cli_sso_attribution_metadata(result=result) if attribution_metadata: await _persist_cli_sso_user_metadata( @@ -2223,8 +2299,8 @@ async def _complete_cli_sso_callback_session( "user_role": user_info.user_role, "models": user_info.models if hasattr(user_info, "models") else [], "user_email": user_email, - "teams": teams, - "team_details": team_details, + "teams": resolved_teams, + "team_details": [detail.model_dump() for detail in team_details], "attribution_metadata": attribution_metadata, } flow["sso_complete"] = True @@ -2233,7 +2309,10 @@ async def _complete_cli_sso_callback_session( _set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow) verbose_proxy_logger.info( - "Stored CLI SSO session for user: %s, teams: %s, num_teams: %s", user_info.user_id, teams, len(teams) + "Stored CLI SSO session for user: %s, teams: %s, num_teams: %s", + user_info.user_id, + resolved_teams, + len(resolved_teams), ) verify_url: Final = get_custom_url( request_base_url=str(request.base_url), @@ -2401,11 +2480,14 @@ async def cli_poll_key( # If no team_id provided and user has 0 or 1 team, use first team (or None) team_id = user_teams[0] if len(user_teams) > 0 else None - team_alias = None - if team_id and isinstance(user_team_details, list): - team_alias = next( - (team.get("team_alias") for team in user_team_details if team.get("team_id") == team_id), - None, + selected_team: Final = _selected_cli_sso_team_detail( + team_details=user_team_details, + team_id=team_id, + ) + if selected_team is None: + raise HTTPException( + status_code=500, + detail=f"Could not resolve the model grants for team: {team_id}. Please run `lite login` again", ) user_info: Final = LiteLLM_UserTable( @@ -2417,7 +2499,9 @@ async def cli_poll_key( jwt_token: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token( user_info=user_info, team_id=team_id, - team_alias=team_alias, + team_alias=selected_team.team_alias, + team_models=selected_team.team_models, + team_model_aliases=selected_team.team_model_aliases, max_budget=None, ) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index a4908af561a..7efd32288e4 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -57,9 +57,13 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable.model_validate(data) - async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member]: + async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member] | None: """Return the team's members_with_roles, locking the row FOR UPDATE. + ``None`` when the team row is gone, which a caller holding the lock can + only see if a delete committed under it, as opposed to ``[]`` for a team + that simply has no members. + Must be called inside a transaction so the row lock is held until commit. This serializes concurrent membership writers on the team row so the losing writer appends onto the winner's committed result instead @@ -69,7 +73,9 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE', team_id, ) - raw_value: Final = rows[0]["members_with_roles"] if rows else None + if not rows: + return None + raw_value: Final = rows[0]["members_with_roles"] parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0377996021c..4892e3b348c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -316,6 +316,9 @@ class LiteLLMCompletionResponsesConfig: "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, } + if not tools: + litellm_completion_request.pop("tool_choice", None) + litellm_completion_request.pop("tools", None) # Responses API `Completed` events require usage, we pass `stream_options` to litellm.completion to include usage if stream is True: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 12c3c5e2be9..b288269b0a2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6164,7 +6164,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4": { "cache_read_input_token_cost": 2.8e-07, @@ -6199,7 +6202,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4": { "cache_read_input_token_cost": 2.8e-07, @@ -6234,7 +6240,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -6276,7 +6285,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, @@ -6312,7 +6324,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, @@ -6348,7 +6363,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -7301,8 +7319,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, @@ -7337,8 +7355,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano": { "cache_read_input_token_cost": 2e-08, @@ -7372,8 +7390,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano-2026-03-17": { "cache_read_input_token_cost": 2e-08, @@ -7408,8 +7426,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_none_reasoning_effort": false, - "supports_xhigh_reasoning_effort": false + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true }, "azure/gpt-image-1": { "cache_read_input_token_cost": 1.25e-06, @@ -8712,6 +8730,268 @@ "/v1/images/generations" ] }, + "azure_ai/FW-DeepSeek-V3.2": { + "cache_read_input_token_cost": 3.1e-07, + "input_cost_per_token": 6.2e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1.85e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-DeepSeek-V4-Pro": { + "cache_read_input_token_cost": 1.65e-07, + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.828e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5": { + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 1.1e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.52e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.1": { + "cache_read_input_token_cost": 2.86e-07, + "input_cost_per_token": 1.54e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 202800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.2": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.54e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.2-Fast": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-Inkling": { + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://fireworks.ai/models/fireworks/inkling", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-Kimi-K2.5": { + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6.6e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.3e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Kimi-K2.6": { + "cache_read_input_token_cost": 1.76e-07, + "input_cost_per_token": 1.045e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Kimi-K2.7-Code": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 1.05e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Kimi-K3": { + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-kimi-k3-through-fireworks-ai-on-microsoft-foundry/4540187", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-MiniMax-M2.5": { + "cache_read_input_token_cost": 3.3e-08, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-MiniMax-M3": { + "cache_read_input_token_cost": 6.6e-08, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "max_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/fireworks/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/FW-Nemotron-3-Ultra-NVFP4": { + "cache_read_input_token_cost": 1.19e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://fireworks.ai/models/fireworks/nemotron-3-ultra-nvfp4", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/MAI-Image-2.5": { "input_cost_per_image_token": 8e-06, "input_cost_per_token": 5e-06, @@ -9329,6 +9609,24 @@ "supports_tool_choice": true, "supports_web_search": true }, + "azure_ai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-grok-4-3-on-microsoft-foundry-latest-generation-agentic-capabilities/4517096", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, @@ -24869,7 +25167,10 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true }, "gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -27711,6 +28012,93 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.2": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.2-contributor": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1e-07, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index ccf710c5708..93c6cfc42d0 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -163,6 +163,7 @@ async def test_team_object_has_object_permission_id(): token=hashed_key, last_refreshed_at=time.time(), team_object_permission_id=permission_id, + team_models=["gpt-4o"], ) user_api_key_cache.set_cache(key=hashed_key, value=valid_token) @@ -255,6 +256,7 @@ async def test_aaauser_personal_budgets(key_ownership): user_id=_user_id, team_id="my-special-team", team_max_budget=100, + team_models=["gpt-4o"], spend=20, ) diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 2e75039139c..a541ab2b3c6 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1,7 +1,7 @@ import json import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -110,15 +110,19 @@ def test_azure_ai_grok_stop_parameter_handling(): config = AzureAIStudioConfig() # Test Grok model detection - assert config._supports_stop_reason("grok-4-fast") == False - assert config._supports_stop_reason("grok-4") == False - assert config._supports_stop_reason("grok-3-mini") == False - assert config._supports_stop_reason("grok-code-fast") == False - assert config._supports_stop_reason("gpt-4") == True + assert config._supports_stop_reason("grok-4-fast") is False + assert config._supports_stop_reason("grok-4.3") is False + assert config._supports_stop_reason("grok-4") is False + assert config._supports_stop_reason("grok-3-mini") is False + assert config._supports_stop_reason("grok-code-fast") is False + assert config._supports_stop_reason("gpt-4") is True # Test supported parameters for Grok models - grok_params = config.get_supported_openai_params("grok-4-fast") - assert "stop" not in grok_params, "Grok models should not support stop parameter" + for model in ("grok-4-fast", "grok-4.3"): + grok_params = config.get_supported_openai_params(model) + assert ( + "stop" not in grok_params + ), "Grok models should not support stop parameter" # Test supported parameters for non-Grok models gpt_params = config.get_supported_openai_params("gpt-4") diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py new file mode 100644 index 00000000000..9917ab41b42 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py @@ -0,0 +1,204 @@ +""" +Regression tests for Azure AI Foundry Fireworks (FW-*) model cost map entries. + +Prices for Data Zone pay-per-token meters come from the Azure retail prices API +(product "Azure Fireworks Models"). Kimi K3 rates come from the Microsoft Foundry +announcement. Models without dedicated Azure meters use published Fireworks +serverless rates. +""" + +import json +from importlib.resources import files + +import pytest + +FW_MODELS = { + "azure_ai/FW-Kimi-K2.5": { + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 1.1e-07, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_vision": True, + }, + "azure_ai/FW-Kimi-K2.6": { + "input_cost_per_token": 1.045e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 1.76e-07, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_vision": True, + }, + "azure_ai/FW-Kimi-K2.7-Code": { + "input_cost_per_token": 1.05e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.1e-07, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_vision": True, + }, + "azure_ai/FW-Kimi-K3": { + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "cache_read_input_token_cost": 3.3e-07, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "supports_vision": True, + }, + "azure_ai/FW-Inkling": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4.05e-06, + "cache_read_input_token_cost": 1.7e-07, + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + }, + "azure_ai/FW-DeepSeek-V3.2": { + "input_cost_per_token": 6.2e-07, + "output_cost_per_token": 1.85e-06, + "cache_read_input_token_cost": 3.1e-07, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + }, + "azure_ai/FW-DeepSeek-V4-Pro": { + "input_cost_per_token": 1.925e-06, + "output_cost_per_token": 3.828e-06, + "cache_read_input_token_cost": 1.65e-07, + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + }, + "azure_ai/FW-MiniMax-M3": { + "input_cost_per_token": 3.3e-07, + "output_cost_per_token": 1.32e-06, + "cache_read_input_token_cost": 6.6e-08, + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "supports_vision": True, + }, + "azure_ai/FW-MiniMax-M2.5": { + "input_cost_per_token": 3.3e-07, + "output_cost_per_token": 1.32e-06, + "cache_read_input_token_cost": 3.3e-08, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + }, + "azure_ai/FW-Nemotron-3-Ultra-NVFP4": { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1.19e-07, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + }, + "azure_ai/FW-GLM-5.2-Fast": { + "input_cost_per_token": 2.1e-06, + "output_cost_per_token": 6.6e-06, + "cache_read_input_token_cost": 2.1e-07, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + }, + "azure_ai/FW-GLM-5.2": { + "input_cost_per_token": 1.54e-06, + "output_cost_per_token": 4.84e-06, + "cache_read_input_token_cost": 1.5e-07, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + }, + "azure_ai/FW-GLM-5.1": { + "input_cost_per_token": 1.54e-06, + "output_cost_per_token": 4.84e-06, + "cache_read_input_token_cost": 2.86e-07, + "max_input_tokens": 202800, + "max_output_tokens": 131072, + }, + "azure_ai/FW-GLM-5": { + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 3.52e-06, + "cache_read_input_token_cost": 2.2e-07, + "max_input_tokens": 200000, + "max_output_tokens": 128000, + }, +} + + +@pytest.fixture(scope="module") +def use_local_model_cost_map(): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + import litellm + from litellm.utils import _invalidate_model_cost_lowercase_map + + original_model_cost = litellm.model_cost + litellm.model_cost = json.loads( + files("litellm") + .joinpath("model_prices_and_context_window_backup.json") + .read_text(encoding="utf-8") + ) + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + try: + yield litellm + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + monkeypatch.undo() + + +@pytest.mark.parametrize("model_key,expected", list(FW_MODELS.items())) +def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected): + model_info = use_local_model_cost_map.get_model_info(model=model_key) + + assert model_info["litellm_provider"] == "azure_ai" + assert model_info["mode"] == "chat" + assert model_info["input_cost_per_token"] == pytest.approx(expected["input_cost_per_token"]) + assert model_info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"]) + assert model_info["cache_read_input_token_cost"] == pytest.approx( + expected["cache_read_input_token_cost"] + ) + assert model_info["max_input_tokens"] == expected["max_input_tokens"] + assert model_info["max_output_tokens"] == expected["max_output_tokens"] + assert model_info["max_tokens"] == expected["max_output_tokens"] + assert model_info["supports_function_calling"] is True + assert model_info["supports_reasoning"] is True + assert model_info["supports_tool_choice"] is True + assert model_info["supports_prompt_caching"] is True + if expected.get("supports_vision"): + assert model_info["supports_vision"] is True + + +@pytest.mark.parametrize( + "model_name,expected_prompt,expected_completion", + [ + ("FW-Kimi-K2.6", 1.045, 4.4), + ("FW-DeepSeek-V4-Pro", 1.925, 3.828), + ("FW-GLM-5.2", 1.54, 4.84), + ("FW-Kimi-K3", 3.3, 16.5), + ("FW-MiniMax-M2.5", 0.33, 1.32), + ("FW-Inkling", 1.0, 4.05), + ("FW-Nemotron-3-Ultra-NVFP4", 0.6, 2.4), + ], +) +def test_azure_ai_fw_cost_per_token( + use_local_model_cost_map, model_name, expected_prompt, expected_completion +): + from litellm.llms.azure_ai.cost_calculator import cost_per_token + from litellm.types.utils import Usage + + usage = Usage( + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + total_tokens=2_000_000, + ) + + prompt_cost, completion_cost = cost_per_token(model=model_name, usage=usage) + + assert prompt_cost == pytest.approx(expected_prompt) + assert completion_cost == pytest.approx(expected_completion) + + +def test_azure_ai_fw_kimi_k26_case_insensitive_lookup(use_local_model_cost_map): + upper = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Kimi-K2.6") + lower = use_local_model_cost_map.get_model_info(model="azure_ai/fw-kimi-k2.6") + + assert upper["input_cost_per_token"] == pytest.approx(lower["input_cost_per_token"]) + assert upper["output_cost_per_token"] == pytest.approx(lower["output_cost_per_token"]) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 298f8a31b64..3ed4c9e9a6d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -155,6 +155,55 @@ def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_value assert token_data["team_alias"] == "test-team" +def test_get_cli_jwt_auth_token_carries_team_grants_not_user_allowlist( + valid_sso_user_defined_values, +): + """A team-bound `lite login` session token must snapshot the team's grants. + + Without team_models the /v1/models bail-out (`not key_models and not team_models`) + treats the session as unrestricted and lists the whole proxy; without + team_model_aliases a team alias never resolves on /chat/completions. The user's + personal allowlist must stay out of the key `models` slot, since a team-bound + credential is governed by the team grant, not by a per-user list. + """ + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + valid_sso_user_defined_values, + team_id="team-123", + team_alias="test-team", + team_models=("claude-sonnet-4-5", "gpt-4.1"), + team_model_aliases={"team-fast": "gpt-4.1-mini"}, + ) + + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data["team_id"] == "team-123" + assert token_data["team_models"] == ["claude-sonnet-4-5", "gpt-4.1"] + assert token_data["team_model_aliases"] == {"team-fast": "gpt-4.1-mini"} + assert valid_sso_user_defined_values.models == ["gpt-3.5-turbo"] + assert token_data["models"] == [] + + +def test_get_cli_jwt_auth_token_keeps_user_allowlist_when_no_team( + valid_sso_user_defined_values, +): + """A session token with no team bound still carries the user's own allowlist.""" + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data.get("team_id") is None + assert token_data["models"] == ["gpt-3.5-turbo"] + assert token_data["team_models"] == [] + + def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( valid_sso_user_defined_values, ): @@ -2073,6 +2122,53 @@ async def test_get_team_object_raises_404_when_not_found(): assert "Team doesn't exist in db" in str(exc_info.value.detail) +def _mock_prisma_for_team_lookup(find_unique): + from unittest.mock import MagicMock + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = find_unique + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_get_team_object_distinguishes_absent_team_from_unreadable_row(): + """A deleted team and a database that would not answer both surface as a 404, + which leaves callers unable to tell a definitive answer from a degraded read. + Only the row being positively absent raises the subclass; anything else keeps + the plain 404 so every existing caller is unaffected.""" + from unittest.mock import AsyncMock, MagicMock + + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import TeamNotFoundError, get_team_object + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + # The database answered, and the row is not there. + with pytest.raises(TeamNotFoundError) as absent_info: + await get_team_object( + team_id="absent-team-lit5522", + prisma_client=_mock_prisma_for_team_lookup(AsyncMock(return_value=None)), + user_api_key_cache=mock_cache, + check_db_only=True, + ) + assert absent_info.value.status_code == 404 + assert "Team doesn't exist in db" in str(absent_info.value.detail) + + # The database did not answer. Same status and detail, but not the subclass, + # so a caller keying on it does not read this as proof the team is gone. + with pytest.raises(HTTPException) as unreadable_info: + await get_team_object( + team_id="unreadable-team-lit5522", + prisma_client=_mock_prisma_for_team_lookup(AsyncMock(side_effect=ConnectionError("db unreachable"))), + user_api_key_cache=mock_cache, + check_db_only=True, + ) + assert unreadable_info.value.status_code == 404 + assert not isinstance(unreadable_info.value, TeamNotFoundError) + + # Reject Client-Side Metadata Tags Tests diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 129813d806c..eea556b9a0d 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4368,6 +4368,212 @@ async def test_centralized_common_checks_team_404_does_not_zero_other_contexts() setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_centralized_common_checks_unresolvable_team_without_grant_is_refused(): + """The store restricts the team to gpt-4o-mini and the read of it fails, so the + only surviving team record is the token's own, which carries ``team_models=[]`` + and reads as every model. The request must be refused with the original lookup + error. Pre-fix it was served.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import HTTPException, Request + from starlette.datastructures import URL + + # The key inherits its models from the team (models=[]), so the team object + # is the only gate on model access. + token = UserAPIKeyAuth( + api_key="sk-test", + team_id="restricted-team", + models=[], + team_models=[], + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"model": "gpt-4.1"}).encode() + + team_read_failure = HTTPException( + status_code=404, + detail={"error": "Team doesn't exist in db. Team=restricted-team."}, + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=team_read_failure, + ): + with pytest.raises(HTTPException) as exc_info: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4.1"}, + route="/chat/completions", + ) + assert exc_info.value is team_read_failure + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_team_models", [[], ["gpt-4.1"]]) +async def test_centralized_common_checks_absent_team_refused_despite_db_unavailable_optout(token_team_models): + """A team that is provably gone is a definitive answer, not a degraded read. + ``allow_requests_on_db_unavailable`` is a static settings read, so without the + absent-versus-unreadable distinction it would hand a deleted team's key the + old permissive fallback while the database is perfectly healthy. Refused in + both token shapes, including the one whose grant would otherwise vouch. + + Imported from the module under test rather than from ``auth_checks``: other + tests in this suite ``importlib.reload`` that module, which rebinds the class + and would leave this raising a type the guard has never seen.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import HTTPException, Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import TeamNotFoundError + + token = UserAPIKeyAuth( + api_key="sk-test", + team_id="deleted-team", + models=[], + team_models=token_team_models, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"model": "gpt-4.1"}).encode() + + team_absent = TeamNotFoundError(team_id="deleted-team") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["general_settings"] = {"allow_requests_on_db_unavailable": True} + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=team_absent, + ): + with pytest.raises(HTTPException) as exc_info: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4.1"}, + route="/chat/completions", + ) + assert exc_info.value is team_absent + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_unreadable_team_keeps_db_unavailable_optout(): + """The counterpart: an unreadable team leaves the grant unknown rather than + answered, so an operator who has accepted degraded authorization during a + database fault still gets the fallback. Without this the fix would trade the + widening for a lockout with no way out.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import HTTPException as _HTTPException + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", team_id="unreadable-team", models=[], team_models=[]) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"model": "gpt-4.1"}).encode() + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["general_settings"] = {"allow_requests_on_db_unavailable": True} + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=_HTTPException(status_code=404, detail={"error": "team unreadable"}), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4.1"}, + route="/chat/completions", + ) + mock_checks.assert_awaited_once() + assert mock_checks.call_args.kwargs["team_object"].team_id == "unreadable-team" + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model, is_granted", + [("gpt-4o-mini", True), ("gpt-4.1", False)], +) +async def test_centralized_common_checks_unresolvable_team_with_grant_enforces_it(requested_model, is_granted): + """Mirror of the refusal above: a token that does carry a team model grant keeps + the fallback, and the reconstructed team must still enforce that grant rather + than wave the request through.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import HTTPException, Request + from starlette.datastructures import URL + + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + token = UserAPIKeyAuth( + api_key="sk-test", + team_id="restricted-team", + models=[], + team_models=["gpt-4o-mini"], + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"model": requested_model}).encode() + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=404, detail={"error": "team unreadable"}), + ): + if is_granted: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": requested_model}, + route="/chat/completions", + ) + else: + with pytest.raises(ProxyException) as exc_info: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": requested_model}, + route="/chat/completions", + ) + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_centralized_common_checks_user_http_exception_isolates_to_user_only(): """Per-fetch isolation, mirror of the team case: an HTTPException diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index afd1696a89f..a23c573047f 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -1,3 +1,4 @@ +import inspect import os import sys from unittest.mock import patch @@ -14,6 +15,9 @@ sys.path.insert( from litellm.proxy.client.cli.commands.agents import ( AgentRunError, + _hand_off, + _replace_process, + _spawn_and_wait, agent_commands, agent_launch_args, agent_profile, @@ -29,11 +33,25 @@ def _agent_command(name): return next(c for c in agent_commands() if c.name == name) +def _default_of(func, param): + return inspect.signature(func).parameters[param].default + + class _FakeResponse: def __init__(self, status_code): self.status_code = status_code +class _Recorder: + def __init__(self, returns=None): + self.returns = returns + self.calls = [] + + def __call__(self, *args): + self.calls.append(args) + return self.returns + + class TestAgentProfile: def test_claude_is_anthropic(self): name, profiles = agent_profile("claude") @@ -314,6 +332,267 @@ class TestRunAgent: assert order == ["launch"] +_WINDOWS_CLAUDE_EXE = "C:\\Program Files\\Claude\\claude.exe" +_WINDOWS_CLAUDE_CMD = "C:\\Users\\dev\\AppData\\Roaming\\npm\\claude.cmd" +_AGENT_ENV = {"ANTHROPIC_BASE_URL": "http://localhost:4000"} +_CMD_PREFIX = "cmd.exe /d /e:on /v:off /s /c " + + +def _shim_command_line(*args): + spawn = _Recorder(returns=0) + with pytest.raises(SystemExit): + _hand_off( + _WINDOWS_CLAUDE_CMD, + ["claude", *args], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=spawn, + ) + return spawn.calls[0][0] + + +class TestHandOff: + def test_windows_spawns_child_instead_of_exec(self): + replace = _Recorder() + spawn = _Recorder(returns=0) + + with pytest.raises(SystemExit) as excinfo: + _hand_off( + _WINDOWS_CLAUDE_EXE, + ["claude", "--resume"], + _AGENT_ENV, + platform="win32", + replace=replace, + spawn=spawn, + ) + + assert excinfo.value.code == 0 + assert replace.calls == [] + assert spawn.calls == [ + ((_WINDOWS_CLAUDE_EXE, "--resume"), _AGENT_ENV), + ] + + @pytest.mark.parametrize("code", [1, 42, 130]) + def test_windows_propagates_child_exit_code(self, code): + with pytest.raises(SystemExit) as excinfo: + _hand_off( + _WINDOWS_CLAUDE_EXE, + ["claude"], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=_Recorder(returns=code), + ) + assert excinfo.value.code == code + + @pytest.mark.parametrize( + "path", + [ + _WINDOWS_CLAUDE_CMD, + "C:\\shims\\claude.CMD", + "C:\\shims\\claude.bat", + ], + ) + def test_windows_batch_shim_goes_through_cmd_exe(self, path): + spawn = _Recorder(returns=0) + + with pytest.raises(SystemExit): + _hand_off( + path, + ["claude", "--resume"], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=spawn, + ) + + assert spawn.calls[0][0] == f'{_CMD_PREFIX}""{path}" "--resume""' + + def test_windows_shim_quotes_a_path_containing_spaces(self): + spawn = _Recorder(returns=0) + path = "C:\\Program Files\\npm\\claude.cmd" + + with pytest.raises(SystemExit): + _hand_off( + path, + ["claude", "-p", "hello world"], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=spawn, + ) + + expected = f'{_CMD_PREFIX}""C:\\Program Files\\npm\\claude.cmd" "-p" "hello world""' + assert spawn.calls[0][0] == expected + + @pytest.mark.parametrize("payload", ["a&calc", "a|calc", "a>out", "a^b", "a&&calc"]) + def test_windows_shim_never_leaves_a_metacharacter_unquoted(self, payload): + expected = f'{_CMD_PREFIX}""{_WINDOWS_CLAUDE_CMD}" "-p" "{payload}""' + assert _shim_command_line("-p", payload) == expected + + def test_windows_shim_doubles_an_embedded_quote(self): + assert _shim_command_line("-p", 'say "hi"').endswith('"-p" "say ""hi""""') + + @pytest.mark.parametrize( + "payload, quoted", + [ + ("%PATH%", "%%cd:~,%PATH%%cd:~,%"), + ("100%", "100%%cd:~,%"), + ("%OS%%CD%", "%%cd:~,%OS%%cd:~,%%%cd:~,%CD%%cd:~,%"), + ], + ) + def test_windows_shim_stops_cmd_expanding_a_percent_variable(self, payload, quoted): + assert _shim_command_line("-p", payload).endswith(f'"-p" "{quoted}""') + + def test_windows_shim_guards_a_percent_in_the_shim_path(self): + spawn = _Recorder(returns=0) + path = "C:\\dev%HOME%\\claude.cmd" + + with pytest.raises(SystemExit): + _hand_off( + path, + ["claude"], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=spawn, + ) + + assert spawn.calls[0][0] == f'{_CMD_PREFIX}""C:\\dev%%cd:~,%HOME%%cd:~,%\\claude.cmd""' + + @pytest.mark.parametrize( + "payload, quoted", + [ + ("C:\\dir\\", "C:\\dir\\\\"), + ('say \\"hi', 'say \\\\""hi'), + ('a\\\\"b', 'a\\\\\\\\""b'), + ], + ) + def test_windows_shim_doubles_backslashes_that_precede_a_quote(self, payload, quoted): + assert _shim_command_line("-p", payload).endswith(f'"-p" "{quoted}""') + + @pytest.mark.parametrize("payload", ["one\ntwo", "one\r\ntwo", "trailing\r"]) + def test_windows_shim_refuses_an_argument_holding_a_line_break(self, payload): + with pytest.raises(AgentRunError, match="line break"): + _hand_off( + _WINDOWS_CLAUDE_CMD, + ["claude", "-p", payload], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=_Recorder(returns=0), + ) + + def test_windows_shim_keeps_the_switches_the_quoting_depends_on(self): + command = _shim_command_line("-p", "hi") + assert command.startswith("cmd.exe ") + switches = command.split(" /c ")[0].split()[1:] + assert switches == ["/d", "/e:on", "/v:off", "/s"] + + def test_windows_exe_is_not_wrapped_in_cmd_exe(self): + spawn = _Recorder(returns=0) + with pytest.raises(SystemExit): + _hand_off( + _WINDOWS_CLAUDE_EXE, + ["claude"], + _AGENT_ENV, + platform="win32", + replace=_Recorder(), + spawn=spawn, + ) + assert spawn.calls[0][0] == (_WINDOWS_CLAUDE_EXE,) + + @pytest.mark.parametrize("platform", ["darwin", "linux", "freebsd8"]) + def test_posix_still_replaces_the_process(self, platform): + replace = _Recorder() + spawn = _Recorder(returns=0) + + _hand_off( + "/usr/local/bin/claude", + ["claude", "--resume"], + _AGENT_ENV, + platform=platform, + replace=replace, + spawn=spawn, + ) + + assert spawn.calls == [] + assert replace.calls == [ + ("/usr/local/bin/claude", ["claude", "--resume"], _AGENT_ENV), + ] + path, args, env = replace.calls[0] + assert isinstance(args, list) + assert isinstance(env, dict) + + def test_replace_process_calls_execvpe_with_argv_and_env(self): + execvpe = _Recorder() + + _replace_process( + "/usr/local/bin/claude", + ("claude", "--resume"), + _AGENT_ENV, + execvpe=execvpe, + ) + + assert execvpe.calls == [ + ("/usr/local/bin/claude", ["claude", "--resume"], _AGENT_ENV), + ] + _path, argv, env = execvpe.calls[0] + assert isinstance(argv, list) + assert isinstance(env, dict) + + def test_posix_default_replacement_is_execvpe(self): + assert _default_of(run_agent, "launcher") is _hand_off + assert _default_of(_hand_off, "replace") is _replace_process + assert _default_of(_replace_process, "execvpe") is os.execvpe + assert _default_of(_hand_off, "spawn") is _spawn_and_wait + assert _default_of(_hand_off, "platform") == sys.platform + + def test_spawn_and_wait_blocks_until_the_child_is_done(self, tmp_path): + marker = tmp_path / "child-finished" + script = ( + "import os, pathlib, time; time.sleep(0.5); " + "pathlib.Path(os.environ['MARKER']).write_text('done'); " + "raise SystemExit(int(os.environ['RC']))" + ) + + code = _spawn_and_wait( + [sys.executable, "-c", script], + {"RC": "7", "MARKER": str(marker), "PATH": os.environ.get("PATH", "")}, + ) + + assert marker.read_text() == "done" + assert code == 7 + + def test_windows_run_agent_spawns_resolved_binary_with_proxy_args(self): + spawn = _Recorder(returns=3) + replace = _Recorder() + + def launcher(path, args, env): + _hand_off(path, args, env, platform="win32", replace=replace, spawn=spawn) + + with pytest.raises(SystemExit) as excinfo: + run_agent( + "http://localhost:4000", + "sk-key", + ["codex", "exec", "do a thing"], + skip_verify=True, + base_env={}, + which=lambda name: _WINDOWS_CLAUDE_CMD.replace("claude", "codex"), + launcher=launcher, + ) + + assert excinfo.value.code == 3 + assert replace.calls == [] + command, env = spawn.calls[0] + shim = _WINDOWS_CLAUDE_CMD.replace("claude", "codex") + assert command.startswith(f'{_CMD_PREFIX}""{shim}" ') + assert command.endswith('"exec" "do a thing""') + assert '"model_provider=""litellm"""' in command + assert env["OPENAI_API_KEY"] == "sk-key" + + class TestAgentCommands: def setup_method(self): self.runner = CliRunner() @@ -423,6 +702,15 @@ class TestAgentCommands: assert captured["api_key"] == "sk-after-login" mock_get.assert_called_once_with(expected_base_url="http://localhost:4000") + def test_child_exit_code_reaches_the_shell(self): + with patch(f"{AGENTS_MODULE}.run_agent", side_effect=SystemExit(42)): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 42 + def test_agent_run_error_becomes_click_error(self): with patch( f"{AGENTS_MODULE}.run_agent", diff --git a/tests/test_litellm/proxy/client/cli/test_config_commands.py b/tests/test_litellm/proxy/client/cli/test_config_commands.py index 698d6188768..d81ee6bd2b1 100644 --- a/tests/test_litellm/proxy/client/cli/test_config_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_config_commands.py @@ -3,6 +3,7 @@ import os import stat import sys from pathlib import Path +from unittest.mock import patch import pytest from click.testing import CliRunner @@ -18,6 +19,7 @@ from litellm.proxy.client.cli.commands.config import ( save_config, ) from litellm.proxy.client.cli.commands.private_json import write_private_json +from litellm.proxy.client.cli.interface import show_commands @pytest.fixture @@ -179,6 +181,85 @@ class TestConfigUnset: assert "not set" in result.output.lower() +class TestHiddenCommands: + """`hidden_commands` lets a deployment curate what `lite` advertises. + + Two listings exist and both must honor it: click's own `--help` table and the + hand-rolled block the interactive shell prints. + """ + + def test_nothing_is_hidden_by_default(self, cli_runner, isolated_home): + result = cli_runner.invoke(cli, ["--help"]) + + assert result.exit_code == 0, result.output + assert "codex" in result.output + assert "opencode" in result.output + + def test_configured_commands_drop_out_of_help(self, cli_runner, isolated_home): + assert cli_runner.invoke(cli, ["config", "set", "hidden_commands", "codex,opencode"]).exit_code == 0 + + result = cli_runner.invoke(cli, ["--help"]) + + assert result.exit_code == 0, result.output + assert "claude" in result.output + assert "codex" not in result.output + assert "opencode" not in result.output + + def test_configured_commands_drop_out_of_interactive_listing(self, capsys, isolated_home): + save_config({"hidden_commands": "codex,keys"}) + + show_commands() + listing = capsys.readouterr().out + + assert "claude" in listing + assert "codex" not in listing + assert "keys" not in listing + assert "teams" in listing + + def test_hidden_commands_are_still_invokable(self, cli_runner, isolated_home): + """Hiding is about the listing only; anyone already scripting the command keeps working.""" + save_config({"hidden_commands": "codex"}) + + with patch("litellm.proxy.client.cli.commands.agents.run_agent") as run_agent_mock: + result = cli_runner.invoke( + cli, + ["--base-url", "http://localhost:4000", "--api-key", "sk-key", "codex", "exec", "do a thing"], + ) + + assert result.exit_code == 0, result.output + _base_url, _api_key, command = run_agent_mock.call_args.args + assert list(command) == ["codex", "exec", "do a thing"] + + def test_unset_brings_the_commands_back(self, cli_runner, isolated_home): + assert cli_runner.invoke(cli, ["config", "set", "hidden_commands", "codex"]).exit_code == 0 + assert cli_runner.invoke(cli, ["config", "unset", "hidden_commands"]).exit_code == 0 + + assert "codex" in cli_runner.invoke(cli, ["--help"]).output + + def test_set_normalizes_whitespace_and_ordering(self, cli_runner, isolated_home): + result = cli_runner.invoke(cli, ["config", "set", "hidden_commands", " opencode , codex ,"]) + + assert result.exit_code == 0, result.output + assert json.loads(_config_path(isolated_home).read_text()) == {"hidden_commands": "codex,opencode"} + + @pytest.mark.parametrize("value", ["", " ", ",", " , "]) + def test_set_empty_list_rejected(self, cli_runner, isolated_home, value): + """An empty value would silently hide nothing; point users at `config unset` instead.""" + result = cli_runner.invoke(cli, ["config", "set", "hidden_commands", value]) + + assert result.exit_code != 0 + assert "unset" in result.output + assert not _config_path(isolated_home).exists() + + def test_set_space_separated_list_rejected(self, cli_runner, isolated_home): + """`lite config set hidden_commands "codex opencode"` would hide neither.""" + result = cli_runner.invoke(cli, ["config", "set", "hidden_commands", "codex opencode"]) + + assert result.exit_code != 0 + assert "without spaces" in result.output + assert not _config_path(isolated_home).exists() + + class TestConfigHelpers: def test_get_config_file_path_under_home(self, isolated_home): assert get_config_file_path() == str(isolated_home / ".litellm" / "config.json") diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 9619e6f67c9..acd6b9ec35a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1,10 +1,8 @@ import os import sys import types -import json from contextlib import ExitStack from datetime import datetime - from types import SimpleNamespace from typing import List, Optional from unittest.mock import AsyncMock, MagicMock, create_autospec, patch @@ -1693,6 +1691,23 @@ class TestTemporaryMCPSessionEndpoints: credentials_are_encrypted=True, ) + @pytest.mark.asyncio + async def test_get_all_mcp_servers_hides_drafts_without_hiding_legacy_null_rows(self): + """Drafts are addressable only by their own id and must never appear in a listing, but a + bare inequality would also drop pre-approval-workflow rows, since SQL evaluates + NULL != 'draft' as NULL.""" + from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers + + find_rows = AsyncMock(return_value=[]) + with patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + find_rows, + ): + await get_all_mcp_servers(MagicMock()) + + where = find_rows.await_args.args[1] + assert where == {"OR": [{"approval_status": None}, {"approval_status": {"not": "draft"}}]} + @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_or_404(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -5482,7 +5497,9 @@ def _edit_endpoint_patches(old_record, update_mock): ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(side_effect=old_record) if isinstance(old_record, Exception) else AsyncMock(return_value=old_record), + AsyncMock(side_effect=old_record) + if isinstance(old_record, Exception) + else AsyncMock(return_value=old_record), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", @@ -5891,7 +5908,13 @@ def test_bundled_openapi_registry_parses_and_entries_are_well_formed(): registry_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), - "..", "..", "..", "..", "litellm", "proxy", "openapi_registry.json", + "..", + "..", + "..", + "..", + "litellm", + "proxy", + "openapi_registry.json", ) with open(registry_path) as f: registry = json.load(f) 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 5190df7521e..f038af4c7a1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4,7 +4,7 @@ import os import sys from datetime import datetime, timezone from typing import Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException @@ -39,6 +39,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.team_endpoints import ( GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, + _STRIP_DELETED_TEAM_FROM_USERS_SQL, _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, @@ -1829,6 +1830,63 @@ async def test_add_team_members_reconciles_against_freshly_locked_row(): assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"] +@pytest.mark.asyncio +async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request(): + """ + Regression pin for the /team/member_add vs /team/delete race. + + The user row and membership writes land before the reconcile takes the team + row lock, so a /team/delete that commits in between has already run its own + reference sweep and cannot see them. The empty locked SELECT is the only + signal that happened, and leaving it at that would strand the member on a + deleted team id, which authorization paths that trust `user.teams` would + treat as membership if the id were ever recreated. So the request must sweep + the references it just wrote and fail, not report success. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[]) + tx.litellm_teamtable.update = AsyncMock() + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + prisma_client = MagicMock() + prisma_client.tx = MagicMock(return_value=tx_cm) + prisma_client.db.execute_raw = AsyncMock() + prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._process_team_members", + new=AsyncMock(return_value=([], [])), + ): + with pytest.raises(HTTPException) as exc_info: + await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="team-deleted-mid-add", + member=Member(user_id="bob", role="user"), + ), + complete_team_data=LiteLLM_TeamTable(team_id="team-deleted-mid-add", members_with_roles=[]), + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + assert exc_info.value.status_code == 404 + tx.litellm_teamtable.update.assert_not_awaited() + + assert prisma_client.db.execute_raw.await_args_list == [ + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-deleted-mid-add") + ] + prisma_client.db.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"team_id": {"in": ("team-deleted-mid-add",)}} + ) + + def test_add_new_models_to_team_with_existing_models(): """ Test add_new_models_to_team function with existing models @@ -4133,6 +4191,106 @@ async def test_team_member_delete_cleans_verification_tokens( ) +@pytest.mark.parametrize( + "roster_email", + ["Alice@Example.com", "alice-invited-as@example.com"], + ids=["case_variant_of_the_row_email", "email_the_row_never_carried"], +) +@pytest.mark.parametrize("user_row_exists", [True, False]) +@pytest.mark.asyncio +async def test_team_member_delete_by_email_the_user_row_does_not_carry( + user_row_exists, roster_email, mock_db_client, mock_admin_auth +): + """ + Removing a member addressed by user_email drove its user-row and membership cleanup off that raw + email instead of off the user_id the roster entry already carries, so an email the user row does + not literally hold matched nothing and both cleanups silently no-opped behind a 200. + + Both roster emails here are reachable over plain HTTP. /team/member_add resolves an email to a + user case-insensitively but stores the caller's casing in members_with_roles, which produces the + case variant; it also leaves an unmatched email on the entry when no user row carries it at all, + which produces the second. Both converge on the same lookup, so they are parametrized inputs + rather than separate paths, and each one has to detect the bug on its own. + + The user table below is case-sensitive like Postgres, so only a lookup driven by the resolved + user_id finds the row. The user_row_exists=False leg pins the second half on its own: the + membership row has to go even when no user row is left to resolve it from. + """ + from litellm.proxy._types import TeamMemberDeleteRequest + from litellm.proxy.management_endpoints.team_endpoints import team_member_delete + + test_team_id = "team-del-email-case-123" + test_user_id = "user-del-email-case-123" + user_row_email = "alice@example.com" + + mock_team_row = MagicMock() + mock_team_row.model_dump.return_value = { + "team_id": test_team_id, + "members_with_roles": [ + {"user_id": test_user_id, "user_email": roster_email, "role": "user"} + ], + "team_member_permissions": [], + "metadata": {}, + "models": [], + "spend": 0.0, + } + + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team_row + ) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row) + + mock_user_row = MagicMock() + mock_user_row.user_id = test_user_id + mock_user_row.user_email = user_row_email + mock_user_row.teams = [test_team_id] + + async def find_user_rows(where): + if not user_row_exists: + return [] + user_id_filter = where.get("user_id") + if isinstance(user_id_filter, dict) and test_user_id in user_id_filter.get( + "in", [] + ): + return [mock_user_row] + if where.get("user_email") == user_row_email: + return [mock_user_row] + return [] + + mock_db_client.db.litellm_usertable.find_many = AsyncMock( + side_effect=find_user_rows + ) + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + mock_db_client.db.litellm_teammembership = MagicMock() + mock_db_client.db.litellm_teammembership.delete_many = AsyncMock( + return_value=MagicMock() + ) + + mock_db_client.db.litellm_verificationtoken = MagicMock() + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock( + return_value=MagicMock() + ) + + await team_member_delete( + data=TeamMemberDeleteRequest(team_id=test_team_id, user_email=roster_email), + user_api_key_dict=mock_admin_auth, + ) + + if user_row_exists: + mock_db_client.db.litellm_usertable.update.assert_awaited_once_with( + where={"user_id": test_user_id}, + data={"teams": {"set": []}}, + ) + else: + mock_db_client.db.litellm_usertable.update.assert_not_awaited() + + mock_db_client.db.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"team_id": test_team_id, "user_id": test_user_id} + ) + + @pytest.mark.asyncio async def test_new_team_max_budget_exceeds_user_max_budget(): """ @@ -7138,6 +7296,367 @@ async def test_delete_team_persists_deleted_teams(monkeypatch): assert records[0]["litellm_changed_by"] == "admin-user" +@pytest.mark.asyncio +async def test_delete_team_sweeps_references_outside_members_with_roles(monkeypatch): + """ + Regression pin for LIT-5511: a deleted team stayed visible on user records. + + `delete_team` drove all of its cleanup off `team.members_with_roles`, so a user row that + referenced the team by any other route (`/user/update`, SSO sync, a membership row written + without a matching roster entry) kept the dangling team id forever and `/user/info` kept + listing the deleted team. The roster here is deliberately EMPTY, so nothing the per-member + `team_member_delete` path does can make this test pass. + + Both cache keys `_cache_team_object` writes are asserted in the same delete: the id key feeds + `get_team_object` and the alias key feeds the JWT `team_alias_jwt_field` path, so either one + surviving keeps the deleted team resolvable for auth until its TTL expires. + """ + from litellm.proxy._types import DeleteTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + doomed_team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + cache_state_when_rows_deleted = {} + + async def record_cache_state_then_delete(*args, **kwargs): + if kwargs.get("table_name") == "team": + cache_state_when_rows_deleted["doomed_still_cached"] = ( + fresh_cache.get_cache(key="team_id:team-doomed") is not None + ) + return {"deleted_teams": ["team-doomed"]} + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=doomed_team) + mock_prisma_client.delete_data = AsyncMock(side_effect=record_cache_state_then_delete) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + mock_execute_raw = AsyncMock() + mock_prisma_client.db.execute_raw = mock_execute_raw + mock_membership_delete_many = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = mock_membership_delete_many + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + for cached_team_id, cached_alias in ( + ("team-doomed", "doomed-team"), + ("team-kept", "kept-team"), + ): + cached_obj = LiteLLM_TeamTableCachedObj( + team_id=cached_team_id, team_alias=cached_alias + ) + fresh_cache.set_cache(key=f"team_id:{cached_team_id}", value=cached_obj) + fresh_cache.set_cache(key=f"team_alias:{cached_alias}", value=cached_obj) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # array_remove strips just the deleted id in one statement; a read-filter-write of the whole + # array would drop any team a concurrent /team/member_add appended between read and write + assert "array_remove" in _STRIP_DELETED_TEAM_FROM_USERS_SQL + assert mock_execute_raw.await_args_list == [ + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), + ], "the sweep must run once before the team row is deleted and again after, so a member_add racing the delete cannot leave the reference behind" + + # same two passes: the second one reaps a membership row inserted while the delete was running + assert mock_membership_delete_many.await_args_list == [ + call(where={"team_id": {"in": ("team-doomed",)}}), + call(where={"team_id": {"in": ("team-doomed",)}}), + ] + + assert fresh_cache.get_cache(key="team_id:team-doomed") is None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + assert fresh_cache.get_cache(key="team_id:team-kept") is not None + assert fresh_cache.get_cache(key="team_alias:kept-team") is not None + + # Eviction must run AFTER the rows are gone: both writers of these keys hydrate from the db, + # so evicting first lets a concurrent auth lookup re-cache the still-present team. + assert cache_state_when_rows_deleted["doomed_still_cached"] is True + + +@pytest.mark.asyncio +async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(monkeypatch): + """ + A virtual key scoped to the team is deleted from the db with the team, but auth resolves a + cached key object without re-reading the team, so leaving the cache entry behind lets that key + keep buying access until its TTL expires. Verified live: without this eviction the same key + still returns HTTP 200 on /v1/chat/completions right after /team/delete. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + team_key = LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed") + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[team_key]) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + fresh_cache.set_cache(key="hashed-doomed-key", value=UserAPIKeyAuth(token="hashed-doomed-key", team_id="team-doomed")) + fresh_cache.set_cache(key="hashed-unrelated-key", value=UserAPIKeyAuth(token="hashed-unrelated-key")) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert fresh_cache.get_cache(key="hashed-doomed-key") is None + # a key that had nothing to do with the deleted team must survive + assert fresh_cache.get_cache(key="hashed-unrelated-key") is not None + + +@pytest.mark.asyncio +async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cache(monkeypatch): + """ + The reconcile sweep runs after the team row is committed deleted. If it ran before cache + eviction, a sweep failure would return an error with the team gone from the db but still + served from cache, which is the exact bug this PR exists to fix. + """ + from litellm.proxy._types import DeleteTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + # the first sweep succeeds, the post-delete reconcile sweep blows up + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")]) + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team") + fresh_cache.set_cache(key="team_id:team-doomed", value=cached_obj) + fresh_cache.set_cache(key="team_alias:doomed-team", value=cached_obj) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + with pytest.raises(ConnectionError): + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # the delete committed, so the cache must not still be serving the team + assert fresh_cache.get_cache(key="team_id:team-doomed") is None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + + +@pytest.mark.asyncio +async def test_delete_team_broadcasts_cache_invalidation_to_other_workers(monkeypatch): + """ + Evicting locally only reaches the worker that handled the delete. Without the broadcast, every + other worker keeps serving the deleted team, and the deleted team's keys, out of its own + in-memory cache until the TTL, so both stay usable for auth cluster-wide. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")] + ) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + published = [] + + async def record_publish(cache_key): + published.append(cache_key) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.publish_auth_cache_invalidation", record_publish) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # the deleted key first, then both keys `_cache_team_object` writes: miss the alias one and the + # JWT-by-alias path keeps resolving the team, miss the token and the key still authenticates + assert published == ["hashed-doomed-key", "team_id:team-doomed", "team_alias:doomed-team"] + + +@pytest.mark.asyncio +async def test_delete_team_survives_a_failing_cache_backend(monkeypatch): + """ + Cache eviction runs after the reference sweep has already committed, so a cache backend that + is unreachable must not abort the delete. If it did, `/team/delete` would fail with the team + row still present but its user references and membership rows already gone. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.delete_data = mock_delete_data + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + # a key to evict: its eviction runs after the key rows are already deleted, so it must not + # raise either + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")] + ) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + exploding_logging_obj = MagicMock() + exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( + side_effect=ConnectionError("redis is down") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", exploding_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + result = await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert result == {"deleted_teams": ["team-doomed"]} + mock_delete_data.assert_any_await(team_id_list=["team-doomed"], table_name="team") + assert exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_count > 0 + + @pytest.mark.asyncio async def test_team_member_delete_persists_deleted_keys(monkeypatch): from litellm.proxy._types import TeamMemberDeleteRequest diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 979eb09d7db..da805b864fc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2847,6 +2847,19 @@ class TestCLIKeyRegenerationFlow: "user_code_verified": False, "session_data": None, } + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock( + return_value=[ + MagicMock( + model_dump=lambda team_id=team_id: { + "team_id": team_id, + "team_alias": team_id, + "models": [], + } + ) + for team_id in ("team1", "team2") + ] + ) with ( patch.dict( os.environ, @@ -2859,7 +2872,7 @@ class TestCLIKeyRegenerationFlow: "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", return_value=mock_user_info, ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( @@ -3156,9 +3169,9 @@ class TestCLIKeyRegenerationFlow: "user_role": "internal_user", "teams": ["team-a", "team-b", "team-c"], "team_details": [ - {"team_id": "team-a", "team_alias": "Team A"}, - {"team_id": "team-b", "team_alias": "Team B"}, - {"team_id": "team-c", "team_alias": "Team C"}, + {"team_id": "team-a", "team_alias": "Team A", "team_models": []}, + {"team_id": "team-b", "team_alias": "Team B", "team_models": []}, + {"team_id": "team-c", "team_alias": "Team C", "team_models": []}, ], "models": ["gpt-4"], "user_email": "test@example.com", @@ -3225,6 +3238,243 @@ class TestCLIKeyRegenerationFlow: # Verify session was deleted after JWT generation mock_cache.delete_cache.assert_called_once() + @pytest.mark.asyncio + async def test_fetch_cli_sso_team_details_projects_team_grants(self): + """The cached team detail must carry the team's model grants. + + The projection used to drop everything except team_id/team_alias, so the + minted CLI token had no team_models and no team_model_aliases to snapshot. + The joined alias table is stored JSON-encoded, so it has to be decoded here + too, otherwise alias lookup at request time is a substring match on a string. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _fetch_cli_sso_team_details, + ) + + team_row = MagicMock() + team_row.model_dump.return_value = { + "team_id": "team-a", + "team_alias": "Team A", + "models": ["claude-sonnet-4-5", "gpt-4.1"], + "litellm_model_table": { + "id": 7, + "model_aliases": json.dumps({"team-fast": "gpt-4.1-mini"}), + "created_by": "admin", + "updated_by": "admin", + }, + } + find_many = AsyncMock(return_value=[team_row]) + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_many = find_many + + details = await _fetch_cli_sso_team_details( + prisma_client=prisma_client, teams=["team-a"] + ) + + assert find_many.await_args.kwargs["include"] == {"litellm_model_table": True} + assert [detail.model_dump() for detail in details] == [ + { + "team_id": "team-a", + "team_alias": "Team A", + "team_models": ("claude-sonnet-4-5", "gpt-4.1"), + "team_model_aliases": {"team-fast": "gpt-4.1-mini"}, + } + ] + + @pytest.mark.asyncio + async def test_fetch_cli_sso_team_details_separates_lookup_failure_from_no_teams(self): + """A failed lookup must not look like a team that resolved to nothing. + + Both used to return [], so a database blip was indistinguishable from a real + answer. The callback needs them apart: a blip has to fail the login, while a + real empty answer means the team rows are genuinely gone. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _fetch_cli_sso_team_details, + ) + + failing_client = MagicMock() + failing_client.db.litellm_teamtable.find_many = AsyncMock( + side_effect=Exception("connection reset") + ) + assert ( + await _fetch_cli_sso_team_details( + prisma_client=failing_client, teams=["team-a"] + ) + is None + ) + + empty_client = MagicMock() + empty_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + assert ( + await _fetch_cli_sso_team_details( + prisma_client=empty_client, teams=["team-a"] + ) + == () + ) + + @pytest.mark.asyncio + async def test_cli_poll_key_mints_jwt_with_selected_team_grants(self): + """The selected team's grants must reach the mint, not just its alias.""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + session_data = { + "user_id": "grants-user", + "user_role": "internal_user", + "teams": ["team-a", "team-b"], + "team_details": [ + { + "team_id": "team-a", + "team_alias": "Team A", + "team_models": ["gpt-4.1"], + "team_model_aliases": {"a-fast": "gpt-4.1-mini"}, + }, + { + "team_id": "team-b", + "team_alias": "Team B", + "team_models": ["claude-sonnet-4-5"], + "team_model_aliases": {"b-fast": "claude-haiku-4-5"}, + }, + ], + "models": ["personal-only"], + "user_email": "grants@example.com", + } + mock_cache = MagicMock(redis_cache=None) + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } + + with ( + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="minted-token", + ) as mock_get_jwt, + ): + result = await cli_poll_key( + key_id="cli-session-grants", + team_id="team-b", + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["status"] == "ready" + kwargs = mock_get_jwt.call_args.kwargs + assert kwargs["team_id"] == "team-b" + assert kwargs["team_alias"] == "Team B" + assert kwargs["team_models"] == ("claude-sonnet-4-5",) + assert kwargs["team_model_aliases"] == {"b-fast": "claude-haiku-4-5"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "team_details", + [ + pytest.param(None, id="detail_fetch_failed"), + pytest.param( + [{"team_id": "team-other", "team_models": []}], id="selected_team_absent" + ), + pytest.param( + [{"team_id": "team-a", "team_alias": "Team A"}], + id="legacy_detail_without_grants", + ), + ], + ) + async def test_cli_poll_key_refuses_to_mint_when_team_grants_are_unknown( + self, team_details + ): + """An unknown team grant must never be minted as an empty one. + + get_complete_model_list falls through to the whole proxy model list when both + the key allowlist and the team allowlist are empty, and team-bound tokens carry + an empty key allowlist by design. So minting an unresolved team as empty would + hand a team-bound CLI session every model on the proxy. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + mock_cache = MagicMock(redis_cache=None) + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": { + "user_id": "grants-user", + "user_role": "internal_user", + "teams": ["team-a"], + "team_details": team_details, + "models": ["personal-only"], + "user_email": "grants@example.com", + }, + } + + with ( + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="minted-token", + ) as mock_get_jwt, + ): + with pytest.raises(HTTPException) as exc_info: + await cli_poll_key( + key_id="cli-session-grants", + team_id="team-a", + x_litellm_cli_poll_secret="poll-secret", + ) + + assert exc_info.value.status_code == 500 + assert "team-a" in str(exc_info.value.detail) + mock_get_jwt.assert_not_called() + mock_cache.delete_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_cli_poll_key_mints_teamless_session_without_team_grants(self): + """A user with no team still mints, keeping their personal allowlist in the key slot.""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + mock_cache = MagicMock(redis_cache=None) + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": { + "user_id": "teamless-user", + "user_role": "internal_user", + "teams": [], + "team_details": [], + "models": ["personal-only"], + "user_email": "teamless@example.com", + }, + } + + with ( + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="minted-token", + ) as mock_get_jwt, + ): + result = await cli_poll_key( + key_id="cli-session-teamless", + team_id=None, + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["status"] == "ready" + kwargs = mock_get_jwt.call_args.kwargs + assert kwargs["team_id"] is None + assert kwargs["team_models"] == () + assert kwargs["user_info"].models == ["personal-only"] + @pytest.mark.asyncio async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self): """A user with a configured budget must not get the max_ui_session_budget fallback cap.""" @@ -3302,7 +3552,7 @@ class TestCLIKeyRegenerationFlow: "user_id": "unbudgeted-user", "user_role": "internal_user", "teams": ["team-x"], - "team_details": [{"team_id": "team-x", "team_alias": "Team X"}], + "team_details": [{"team_id": "team-x", "team_alias": "Team X", "team_models": []}], "models": ["gpt-4"], "user_email": "unbudgeted@example.com", } @@ -6539,6 +6789,17 @@ class TestCliSsoAttributionMetadata: return_value=MagicMock(metadata={"auth_provider": "generic"}) ) mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock( + return_value=[ + MagicMock( + model_dump=lambda: { + "team_id": "team1", + "team_alias": "team1", + "models": [], + } + ) + ] + ) with ( patch.dict( @@ -7879,6 +8140,117 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): assert response.status_code == 200 +def _cli_callback_kwargs(flow): + return { + "request": _cli_callback_request(), + "key": "cli-login-id", + "flow": flow, + "result": {"sub": "raw-idp-subject"}, + "parsed_openid_result": { + "user_id": "raw-idp-subject", + "user_email": "u@example.com", + "user_role": None, + }, + "user_defined_values": None, + "prisma_client": MagicMock(), + "user_api_key_cache": MagicMock(), + "cli_sso_session_cache": MagicMock(), + "proxy_logging_obj": MagicMock(), + } + + +def _cli_callback_request(): + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + return mock_request + + +def _cli_callback_user_info(teams): + user_info = MagicMock() + user_info.user_id = "cli-user-id" + user_info.user_role = "internal_user" + user_info.models = ["personal-only"] + user_info.teams = teams + return user_info + + +@pytest.mark.asyncio +async def test_cli_completion_drops_teams_whose_rows_no_longer_exist(): + """A membership pointing at a deleted team must not be offered for selection. + + Deleting an organization removes its team rows but leaves the user's membership + behind. If that dead team still reached the session, it would be auto-selected + for a single-team user, its grants could never resolve, and every future login + would be refused with no way for the user to recover. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _CliSsoTeamDetail, + _complete_cli_sso_callback_session, + ) + + live_detail = _CliSsoTeamDetail( + team_id="team-live", team_alias="Live", team_models=("gpt-4.1",) + ) + flow = {} + with ( + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=_cli_callback_user_info(["team-live", "team-deleted"])), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso._fetch_cli_sso_team_details", + AsyncMock(return_value=(live_detail,)), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.build_cli_sso_attribution_metadata", + return_value={}, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + ): + response = await _complete_cli_sso_callback_session(**_cli_callback_kwargs(flow)) + + assert response.status_code == 200 + assert flow["session_data"]["teams"] == ["team-live"] + assert [d["team_id"] for d in flow["session_data"]["team_details"]] == ["team-live"] + + +@pytest.mark.asyncio +async def test_cli_completion_fails_the_login_when_team_lookup_fails(): + """A lookup failure must fail the login instead of caching a teamless session. + + Silently dropping every team here would hand a team-bound user a session with + their personal allowlist, which is the same "unknown grant treated as a real + grant" bug in a quieter form. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _complete_cli_sso_callback_session, + ) + + flow = {} + with ( + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=_cli_callback_user_info(["team-live"])), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso._fetch_cli_sso_team_details", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _complete_cli_sso_callback_session(**_cli_callback_kwargs(flow)) + + assert exc_info.value.status_code == 500 + assert "session_data" not in flow + + class TestSameOriginReturnPath: """The same-origin relative return_to arm added for the MCP gateway DCR authorize round-trip: only strictly relative paths qualify, so login can never redirect the diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 3f567397e1a..38af52f165c 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -546,12 +546,19 @@ class TestTeamRepository: @pytest.mark.asyncio async def test_get_members_with_roles_locked_missing_row(self, repo): + """None, not [], so a caller can tell a deleted team from an empty one. + + /team/member_add reconciles membership under this lock and has to fail, + and clean up the references it already wrote, when a /team/delete + committed underneath it. An empty list would look like a live team with + no members and it would carry on writing. + """ tx = MagicMock() tx.query_raw = AsyncMock(return_value=[]) members = await repo.get_members_with_roles_locked(tx, "missing") - assert members == [] + assert members is None @pytest.mark.asyncio async def test_create_team_all_fields(self, repo): diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index d3aa8eddf45..fef8c2d1349 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -982,6 +982,18 @@ class TestFunctionCallTransformation: assert result["extra_headers"] == {"X-Test-Header": "test-value"} + def test_drops_tool_choice_when_no_tools(self): + """Chat completions providers reject tool_choice when no tools are present.""" + result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="azure_ai/grok-4.3", + input="who are you?", + responses_api_request={"tool_choice": "auto", "tools": []}, + custom_llm_provider="azure_ai", + ) + + assert "tool_choice" not in result + assert "tools" not in result + def test_function_call_without_call_id_fallback_to_id(self): """Test that function_call items can use 'id' field when 'call_id' is missing""" function_call_item = { diff --git a/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py b/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py new file mode 100644 index 00000000000..22cabfbb0eb --- /dev/null +++ b/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py @@ -0,0 +1,83 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm import get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + +AZURE_AI_GROK_4_3_MODEL = "azure_ai/grok-4.3" +AZURE_AI_GROK_4_3_SOURCE = "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-grok-4-3-on-microsoft-foundry-latest-generation-agentic-capabilities/4517096" + + +def _load_model_cost(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +@pytest.fixture(autouse=True) +def reload_model_costs(): + original_model_cost = litellm.model_cost + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + litellm.model_cost = _load_model_cost(json_path) + get_model_info.cache_clear() + yield + litellm.model_cost = original_model_cost + get_model_info.cache_clear() + + +def test_azure_ai_grok_4_3_model_info(): + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + model_cost = _load_model_cost(json_path) + + info = model_cost.get(AZURE_AI_GROK_4_3_MODEL) + assert ( + info is not None + ), f"{AZURE_AI_GROK_4_3_MODEL} not found in model_prices_and_context_window.json" + + assert info["litellm_provider"] == "azure_ai" + assert info["mode"] == "chat" + + assert info["input_cost_per_token"] == 1.25e-06 + assert info["output_cost_per_token"] == 2.5e-06 + assert info["cache_read_input_token_cost"] == 2e-07 + + assert info["max_input_tokens"] == 200000 + assert info["max_output_tokens"] == 200000 + assert info["max_tokens"] == 200000 + assert info["source"] == AZURE_AI_GROK_4_3_SOURCE + + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_web_search"] is True + + routed_model, provider, _, _ = get_llm_provider(model=AZURE_AI_GROK_4_3_MODEL) + assert routed_model == "grok-4.3" + assert provider == "azure_ai" + + resolved_info = get_model_info(model="grok-4.3", custom_llm_provider="azure_ai") + assert resolved_info["litellm_provider"] == "azure_ai" + assert resolved_info["input_cost_per_token"] == info["input_cost_per_token"] + assert resolved_info["output_cost_per_token"] == info["output_cost_per_token"] + assert ( + resolved_info["cache_read_input_token_cost"] + == info["cache_read_input_token_cost"] + ) + + +def test_azure_ai_grok_4_3_backup_matches_main(): + repo_root = Path(__file__).parents[2] + main_path = repo_root / "model_prices_and_context_window.json" + backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json" + + main_cost = _load_model_cost(main_path) + backup_cost = _load_model_cost(backup_path) + + assert backup_cost.get(AZURE_AI_GROK_4_3_MODEL) == main_cost.get( + AZURE_AI_GROK_4_3_MODEL + ) diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py new file mode 100644 index 00000000000..20aa4b11dcd --- /dev/null +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -0,0 +1,121 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import StandardBuiltInToolCostTracking + +MUSE_SPARK_STANDARD = "meta/muse-spark-1.2" +MUSE_SPARK_CONTRIBUTOR = "meta/muse-spark-1.2-contributor" +WEB_SEARCH_COST_PER_QUERY = 0.0025 + +PRICING = ( + (MUSE_SPARK_STANDARD, 1.25e-06, 1.5e-07, 4.25e-06), + (MUSE_SPARK_CONTRIBUTOR, 1e-07, 2e-09, 2e-07), +) + + +def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> dict: + with open(Path(__file__).parents[2] / filename) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) +def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): + info = _load_cost_map().get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["litellm_provider"] == "meta" + assert info["mode"] == "chat" + + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost + assert info["cache_read_input_token_cost"] == cached_cost + + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 131072 + assert info["max_tokens"] == 131072 + + assert info["supports_function_calling"] is True + assert info["supports_parallel_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_pdf_input"] is True + assert info["supports_web_search"] is True + assert info["supports_minimal_reasoning_effort"] is True + assert info["supports_xhigh_reasoning_effort"] is True + + assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + assert info["supported_modalities"] == ["text", "image", "video"] + assert info["supported_output_modalities"] == ["text"] + + assert info["search_context_cost_per_query"] == { + "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, + "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, + "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, + } + + +@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) +def test_muse_spark_1_2_cost_per_token( + local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float +): + prompt_cost, completion_cost = cost_per_token(model=model, prompt_tokens=1000, completion_tokens=500) + + assert prompt_cost == pytest.approx(1000 * input_cost) + assert completion_cost == pytest.approx(500 * output_cost) + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_2_routes_to_meta_model_api(model: str): + routed_model, provider, _, api_base = get_llm_provider(model=model, api_key="sk-test") + + assert routed_model == model.split("/", 1)[1] + assert provider == "meta" + assert api_base == "https://api.meta.ai/v1" + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_2_web_search_cost_per_query(local_model_cost_map, model: str): + info = litellm.get_model_info(model=model) + + assert StandardBuiltInToolCostTracking.get_cost_for_web_search(model_info=info) == WEB_SEARCH_COST_PER_QUERY + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_2_backup_matches_main(model: str): + """Ensure the bundled model cost map stays in sync with the canonical file.""" + main_cost = _load_cost_map() + backup_cost = _load_cost_map("litellm/model_prices_and_context_window_backup.json") + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_muse_spark_contributor_tier_is_cheaper_than_standard(): + cost_map = _load_cost_map() + standard = cost_map[MUSE_SPARK_STANDARD] + contributor = cost_map[MUSE_SPARK_CONTRIBUTOR] + + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert contributor[field] < standard[field], f"contributor {field} should undercut the standard tier" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index 4a6ff77d99b..482901dfd14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -8,6 +8,7 @@ import { ApiError } from "@/lib/http/client"; vi.mock("./useAutoRouterBenchmarks", () => ({ useAutoRouterBenchmarks: vi.fn() })); vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAutoRouters: vi.fn() })); +vi.mock("./ShadowEvalSection", () => ({ default: () =>
})); import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; @@ -274,6 +275,33 @@ describe("AutoRouterBenchmarksTab", () => { expect(screen.getByText("Last 24 hours")).toBeInTheDocument(); }); + it("shows usage by default and mounts shadow evals only when its sub-tab is selected", () => { + mockHook({ data: response([group()]) }); + renderTab(); + + expect(screen.getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByText("Total estimated savings")).toBeInTheDocument(); + expect(screen.queryByTestId("shadow-eval-section")).not.toBeInTheDocument(); + + fireEvent.click(screen.getByRole("tab", { name: "Shadow Evals" })); + expect(screen.getByRole("tab", { name: "Shadow Evals" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByTestId("shadow-eval-section")).toBeInTheDocument(); + + fireEvent.click(screen.getByRole("tab", { name: "Usage" })); + expect(screen.getByText("Total estimated savings")).toBeInTheDocument(); + expect(screen.getByTestId("shadow-eval-section")).toBeInTheDocument(); + }); + + it("keeps the shadow evals sub-tab reachable while the usage body is in its error state", () => { + mockHook({ error: new ApiError("boom", 500, {}) }); + renderTab(); + + expect(screen.getByText("Auto-router usage is unavailable right now")).toBeInTheDocument(); + + fireEvent.click(screen.getByRole("tab", { name: "Shadow Evals" })); + expect(screen.getByTestId("shadow-eval-section")).toBeInTheDocument(); + }); + it("keeps the window picker reachable while a window has no sessions", () => { mockHook({ data: response([]) }); renderTab(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index 5d4fda765e7..80ddef29c9d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -8,7 +8,7 @@ import { Badge } from "@/components/ui/badge"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; -import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { ApiError } from "@/lib/http/client"; import { formatNumberWithCommas } from "@/utils/dataUtils"; @@ -31,6 +31,7 @@ import { type BucketRow, } from "./autoRouterBenchmarks"; import { usd } from "./costOptimizationUtils"; +import ShadowEvalSection from "./ShadowEvalSection"; import TierTurnsChart from "./TierTurnsChart"; import { useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks"; @@ -268,7 +269,7 @@ interface AutoRouterBenchmarksTabProps { accessToken: string | null; } -const AutoRouterBenchmarksTab: React.FC{emptyResultsText(job, resultsError)}
; + } + return ( + <> ++ Router matched or beat your current model +
++ {pct(results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct)} +
+of {(job.judged_count ?? 0).toLocaleString()} judged responses
++ Shadowing {job.shadow_percentage}% via {job.router_name} +
++ {(job.judged_count ?? 0).toLocaleString()} of {job.max_turns.toLocaleString()} turns judged ·{" "} + {(job.error_count ?? 0).toLocaleString()} errored · {usd(job.judge_spend ?? 0)} judge spend + {active && remaining ? ` · ${remaining}` : ""} +
++ Last failure: {job.last_error} +
+ )} ++ Duplicates a sampled slice of the key's traffic through the auto-router and has an LLM judge compare both + answers blind. The router's answers are never served to users; judge calls bill to the shadowed key. +
+Enter a value from 0.1 to 100
+ )} +Enter a value from 1 to 2000
+ )} ++ Would the auto-router have answered as well as the models you use today? Find out on your real traffic, before + switching anything. +
+Existing evaluations could not be loaded. Refresh the page to retry.
+ )} + + {isPending && error == null &&Loading evaluations...
} + + {showcased.map((job) => ( +{children},
+}));
+
+vi.mock("react-syntax-highlighter/dist/esm/styles/prism", () => ({ coy: {} }));
+
+vi.mock("@/contexts/ChatShellContext", () => ({
+ useChatShell: () => {
+ const history = useChatHistory(null, "metrics-test-user");
+ return {
+ accessToken: "sk-test",
+ userId: "metrics-test-user",
+ userEmail: "tester@example.com",
+ userRole: "Admin",
+ premiumUser: false,
+ selectedMCPServers: [],
+ setSelectedMCPServers: vi.fn(),
+ conversations: history.conversations,
+ activeConversation: history.activeConversation,
+ activeConversationId: history.currentActiveId,
+ storageUnavailable: false,
+ staleId: false,
+ createConversation: history.createConversation,
+ appendMessage: history.appendMessage,
+ updateLastAssistantMessage: history.updateLastAssistantMessage,
+ truncateFromMessage: history.truncateFromMessage,
+ deleteConversation: vi.fn(),
+ renameConversation: vi.fn(),
+ };
+ },
+}));
+
+const ONE_TURN_ARG_COUNT = 25;
+const ON_TIMING_DATA_INDEX = 7;
+const ON_USAGE_DATA_INDEX = 8;
+const ON_TOTAL_LATENCY_INDEX = 24;
+
+async function sendOneMessage(): Promise