From cb041966bf3502c56221ea2605e0b0373dce4cac Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 4 Jun 2026 23:37:20 +0530 Subject: [PATCH 001/133] Litellm oss staging 040626 (#29671) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(azure): apply api_version fallback chain to image edit URL `AzureImageEditConfig.get_complete_url` only read `api_version` from `litellm_params`. When callers configured it via `litellm.api_version` or `AZURE_API_VERSION`, the constructed URL had no `?api-version=` and Azure responded `404 Resource not found`. Apply the same fallback chain the Azure chat path already uses in `common_utils.py`: litellm_params > litellm.api_version > AZURE_API_VERSION env > litellm.AZURE_DEFAULT_API_VERSION Adds 5 unit tests pinning each layer of the chain plus a regression guard for `api_base` that already carries `?api-version=`. * feat(mcp): core sampling and elicitation flow with security hardening - Add sampling_handler.py: full MCP sampling/createMessage flow with model selection (hint-based + priority-based), auth enforcement, budget checks, route restriction gates, and tag policy pre-auth - Add elicitation_handler.py: MCP elicitation/create relay with downstream client capability detection - Wire sampling/elicitation callbacks in mcp_server_manager.py gated behind allow_sampling/allow_elicitation config flags - Add allow_sampling/allow_elicitation fields to MCPServer type - Fix session lock deadlock: skip lock for JSON-RPC response POSTs (elicitation/sampling replies) with truncated-body heuristic - Extend client.py with sampling_callback and elicitation_callback - Security: RouteChecks gate, tag-budget bypass fix, x-forwarded-for spoofing fix, Latin-1 header encoding guard - Add 4 new test modules (model access, priority selection, request builder, tool conversion) + update existing MCP tests * fix(security): run pre-call guardrails before MCP sampling acompletion Without this, an upstream MCP server with allow_sampling enabled could send prompts that bypass every guardrail (content filtering, PII redaction, prompt-injection detection) configured on /chat/completions. - Call proxy_logging_obj.pre_call_hook(call_type='acompletion') before llm_router.acompletion so guardrails fire for sampling sub-calls - Add HTTPException to the re-raise list so guardrail rejections propagate correctly instead of being swallowed as generic errors * feat(bedrock_mantle): add Responses API support (/openai/v1/responses) (#29490) * feat(bedrock_mantle): add Responses API transformation config * test(bedrock_mantle): cover trailing-slash api_base normalization * feat(bedrock_mantle): export BedrockMantleResponsesAPIConfig * feat(bedrock_mantle): register gpt-5.x Responses config (gpt-oss unchanged) * feat(bedrock_mantle): add gpt-5.5/gpt-5.4 Responses price-map entries * refactor(bedrock_mantle): exclude gpt-oss instead of allow-listing gpt-5 for Responses routing Frontier OpenAI models on Bedrock Mantle are Responses-only on /openai/v1/responses; gpt-oss is the legacy family that also speaks chat-completions. Gate by excluding gpt-oss (which keeps its chat-completions emulation) and defaulting everything else to the native Responses config, so future frontier models (gpt-6, etc.) route correctly without a code change. Verified against the live us-east-2 Mantle endpoint: gpt-oss 400s on /openai/v1/responses while gpt-5.5 400s on both standard paths. * test(bedrock_mantle): cover supports_native_websocket opt-out Closes the one uncovered line flagged by codecov on the Responses config. The assertion documents that Mantle Responses has no realtime/websocket transport, so realtime routing must not attempt a socket it cannot serve. * fix(bedrock_mantle): route file_search through emulation instead of forwarding to Mantle BedrockMantleResponsesAPIConfig inherited supports_native_file_search() -> True from OpenAIResponsesAPIConfig but never overrode it. Mantle has no OpenAI vector stores, so a forwarded file_search tool is rejected with a 400 (verified upstream: Tool type 'file_search' is not supported). Opting out, like the existing supports_native_websocket override, routes the tool through LiteLLM's file_search emulation instead. * fix(bedrock_mantle): only route openai.gpt frontier models to Responses The previous gate excluded gpt-oss and routed every other model to the native Responses config. But on Mantle only the OpenAI gpt frontier models (gpt-5.x) are served on /openai/v1/responses; gpt-oss and the non-OpenAI families (nvidia, mistral, google, zai, ...) are chat-completions only and 400 on that path. Allow-list the openai.gpt- family (excluding gpt-oss) instead, so chat-only models fall through to the chat-completions emulation. Verified against the live us-east-2 endpoint: nvidia.nemotron-nano-9b-v2 returns 400 on /openai/v1/responses and 200 on /v1/chat/completions. * feat(custom_llm): allow streaming/astreaming to yield ModelResponseStream (#27580) * fix(custom_llm): allow streaming/astreaming to yield ModelResponseStream directly * fix(streaming): enhance ModelResponseStream handling for custom LLM providers * fix(streaming): strip finish_reason from content chunks and ensure tool_calls are preserved * fix(streaming): add type ignore for finish_reason assignment in CustomStreamWrapper * fix(proxy): strip stack trace from HTTP 503 responses (CWE-209) (#28330) * fix(proxy/cwe-209): strip Python traceback from HTTP 503 error responses The /cache/ping endpoint included a full Python traceback in its 503 error response body (inside the ProxyException message), leaking internal file paths, line numbers, and call stacks to any caller. Two MCP route handlers in proxy_server.py similarly interpolated str(e) into "Internal server error" detail strings. Fix: log the traceback server-side via verbose_proxy_logger.exception() and omit it from the ProxyException payload / HTTPException detail returned to clients. Tests updated to assert no "traceback" keyword or frame paths appear in the 503 body, with a new dedicated regression test. CWE-209: Generation of Error Message Containing Sensitive Information. Co-Authored-By: Claude Sonnet 4.6 * fix(proxy/cwe-209): apply Greptile P2 fixes and add MCP exception-path tests Greptile 4/5 review identified two remaining gaps and Codecov reported 0% coverage on the two MCP handler exception branches: 1. caching_routes.py — str(e) in "Service Unhealthy ({str(e)})" could still leak Redis hostnames/IPs; replaced with static "Service Unhealthy". HTTPException is now re-raised before the generic handler so the "cache not initialized" 503 still reaches callers with its detail. Removed the redundant str(e) arg from verbose_proxy_logger.exception() (exception() already appends the traceback automatically). 2. tests — two new unit tests cover the exception paths in dynamic_mcp_route and toolset_mcp_route that were previously at 0%: - test_dynamic_mcp_route_unexpected_exception_returns_500_without_traceback - test_toolset_mcp_route_unexpected_exception_returns_500_without_traceback All 25 tests pass (9 caching + 16 MCP). CWE-209: Generation of Error Message Containing Sensitive Information. Co-Authored-By: Claude Sonnet 4.6 * test(caching_routes): restore precise assertion in test_cache_ping_no_cache_initialized The assertion was weakened to `"Cache not initialized" in str(data)`, which matches the raw string of the entire response dict and would pass even if the error moved to an unexpected field or changed structure. Restore a targeted check on the parsed response: assert the exact string in the correct field `data["detail"]`, matching FastAPI's HTTPException serialisation format {"detail": ""}. Co-Authored-By: Claude Sonnet 4.6 * test(caching_routes): restore precise assertion and add CWE-209 no-cache path test The assertion in test_cache_ping_no_cache_initialized was weakened to `"Cache not initialized" in str(data)`, which matched against the raw string representation of the entire response dict. This would pass silently even if the error message moved to an unexpected field or the structure changed. Restore a targeted assertion on the parsed field: assert data["detail"] == "Cache not initialized. litellm.cache is None" matching FastAPI's HTTPException serialisation format exactly. Add test_cache_ping_no_cache_does_not_expose_internals to show the code path is still working correctly after the CWE-209 fix: verifies that the HTTPException is re-raised as-is (no traceback, no source paths), and asserts the complete response structure is exactly {"detail": "Cache not initialized. litellm.cache is None"}. Co-Authored-By: Claude Sonnet 4.6 * fix(caching_routes): restore ProxyException envelope for null-cache 503 The except HTTPException: raise guard (added in the CWE-209 fix) caused the null-cache HTTPException to escape as FastAPI's {"detail": "..."} shape instead of the {"error": {...}} ProxyException envelope that callers expect. Move the null-cache guard before the try block and raise ProxyException directly so the response structure is consistent with all other /cache/ping 503s, and the except HTTPException: raise guard is only reachable by unexpected downstream HTTPExceptions. Update the two no-cache tests to assert the correct ProxyException envelope. Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 * Update utils.py (#26609) * feat(pricing): add Snowflake Cortex REST API model pricing (#26612) * feat(pricing): add Snowflake Cortex REST API model pricing ## Summary Adds pricing and context window information for 20+ Snowflake Cortex REST API models to `model_prices_and_context_window.json`. ## What's included - **7 Claude models** (sonnet-4-5, sonnet-4-6, 4-sonnet, 4-opus, haiku-4-5, 3-7-sonnet, 3-5-sonnet) — with prompt caching rates - **4 OpenAI models** (gpt-4.1, gpt-5, gpt-5-mini, gpt-5-nano) — with prompt caching rates - **5 Llama models** (3.1-8b, 3.1-70b, 3.1-405b, 3.3-70b, 4-maverick) - **1 DeepSeek model** (deepseek-r1) - **1 Mistral model** (mistral-large2) - **1 Snowflake model** (snowflake-llama-3.3-70b) - **2 Embedding models** (arctic-embed-l-v2.0, arctic-embed-m-v2.0) Each entry includes `input_cost_per_token`, `output_cost_per_token`, `cache_read_input_token_cost` (where applicable), `max_input_tokens`, `max_output_tokens`, and capability flags (`supports_function_calling`, `supports_vision`, `supports_prompt_caching`, `supports_reasoning`). ## Pricing source All prices are in USD per token, sourced from the official [Snowflake Service Consumption Table](https://www.snowflake.com/legal-files/CreditConsumptionTable.pdf) — Tables 6(b) (REST API with Prompt Caching) and 6(c) (REST API). ## Context The existing `snowflake/` provider has zero model entries in the pricing JSON, which means LiteLLM cannot track costs for Snowflake Cortex calls. This PR fills that gap. ## Related - Existing provider: `litellm/llms/snowflake/` - Cortex REST API docs: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api * Update model_prices_and_context_window.json Fix the JSON parsing error * Update model_prices_and_context_window.json Removed the duplicate entry * fix(utils): copy extra_body before adding unknown params to prevent model config mutation (#29620) Fixes #29615. In add_provider_specific_params_to_optional_params, the line: extra_body = passed_params.pop("extra_body", None) or {} returns the original dict reference when extra_body is non-empty (truthy). Subsequent writes like extra_body[k] = passed_params[k] then mutate the shared model config object held by the router, poisoning /model/info and all subsequent requests for that deployment. The or {} short-circuit creates a new dict only when extra_body is falsy (None or {}), which is why the bug does not reproduce with extra_body: {}. Fix: wrap in dict() so we always work on a fresh shallow copy. * fix(vertex_ai): Bake tool_choice into Gemini CachedContent body to prevent silent drop (#29097) * fix(vertex_ai): bake tool_choice into Gemini CachedContent body to prevent silent drop * address greptile feedback on tool_choice cache test * adds test that uses ToolConfig(functionCallingConfig=FunctionCallingConfig(mode=ANY)) instead of a dict literal, mirroring what map_tool_choice_values actually produce * fix(gemini/veo): move image from parameters into instances[0] (#29501) * fix(gemini/veo): move image from parameters into instances[0] Veo's predictLongRunning schema puts image (and prompt) on the instances element; parameters is for aspectRatio/durationSeconds/etc. The Gemini path was leaving image in params_copy, so it ended up nested under parameters and the API silently ignored it. The Vertex path already builds the instance dict explicitly, so this just aligns the Gemini path with it. Fixes #29498 * address greptile: unconditional pop + BytesIO test - Pop `image` from params_copy unconditionally so it never reaches GeminiVideoGenerationParameters even when None, removing implicit reliance on Pydantic's extra-field-ignore. - Add test_transform_video_create_request_image_filelike_goes_to_instance covering the BytesIO path (_convert_image_to_gemini_format) — round-trips the base64 to confirm encoding. - Add test_transform_video_create_request_image_none_is_dropped covering the new None branch. * fix(huggingface): handle special token text in embedding usage (#29660) * fix(guardrails): recompile ToolPermissionGuardrail rules on update_in_memory_litellm_params (#29655) * fix(guardrails): recompile ToolPermissionGuardrail rules on update_in_memory_litellm_params ToolPermissionGuardrail builds self.rules and the compiled target/pattern maps only in __init__. The base update_in_memory_litellm_params re-sets raw attributes via setattr but never rebuilds those maps, so a guardrail updated in place (PUT /guardrails, or the immediate in-memory sync) keeps enforcing the construction-time rules until it is reinitialized (PATCH path, periodic DB poll, or restart). Extract the compile step into _load_rules and override update_in_memory_litellm_params to rebuild from it (dict- and model-safe), re-normalizing default_action / on_disallowed_action. Mirrors the existing PresidioGuardrail override of the same method. Adds regression tests. Fixes #29592. * fix(guardrails): handle dict params in ToolPermissionGuardrail in-memory update Delegate to super() only for LitellmParams input (the base setattr loop is model-only); apply the raw-dict case inline. Fixes the mypy arg-type error and makes the recompile work when the proxy passes the raw DB dict. * fix(guardrails): preserve tool-permission rules on a partial in-memory update A partial update (e.g. a LitellmParams whose rules field is None) ran through the generic setattr, which set self.rules to None, and the recompile was skipped, leaving the guardrail with no rules. Snapshot the previous rules and restore them when the update carries no rules; an explicit empty list still clears them. Adds a regression test for the rules-absent case. Addresses the Greptile review note on #29655. * fix(bedrock): stop base_model label from stripping tools/tool_choice (#29621) * fix(bedrock): stop base_model label from stripping tools/tool_choice A Router/proxy Bedrock deployment whose model_info.base_model is a friendly label (e.g. claude-haiku-4-5) silently lost tools/tool_choice: the outgoing Converse request was built without toolConfig, so the model behaved as if no tools were provided. Worked in v1.84.0, regressed in v1.85.0, and with drop_params=true it failed silently. Two changes compound into the bug. completion() passed model_info.base_model as the model argument to get_optional_params, so the real Bedrock model id never reached supported-param resolution; and get_supported_openai_params resolved the provider config's params from base_model or model, letting the label fully replace the real model. For Bedrock the label resolves to no tool support, so tools/tool_choice were dropped before transformation. completion() now keeps model as the real deployment model and threads the resolved base_model (kwarg or model_info) through separately, and get_supported_openai_params treats base_model as additive: it returns the union of the params supported by model and by base_model. A hint can only add capabilities, never strip ones the real model already exposes, which also preserves the original base_model behavior from #27717 and Azure's base_model driven model-type detection. Fixes #29618 * test(main): make base_model param test robust to new parametrize cases Restore an explicit per-case expected_model_param literal instead of hardcoding the gemini id, so a future case with a different model can't produce a misleading assertion failure. * fix(fireworks_ai): pass response_format json_schema through unchanged (#29606) FireworksAIConfig.map_openai_params was rewriting the OpenAI strict `{type: json_schema, json_schema: {name, strict, schema}}` shape into `{type: json_object, schema: ...}` before sending to Fireworks, dropping `strict` and `name` and changing the `type`. Per Fireworks' docs json_object means "force any valid JSON output (no specific schema)", so the schema constraint was effectively dropped and grammar-guided decoding never ran; model output silently violated the schema. The rewrite landed in #7085 (Dec 2024) when Fireworks did not yet accept native json_schema. Fireworks accepts the OpenAI strict shape natively now, so the rewrite has become a regression. Removes the rewrite. Passes response_format through unchanged. Updates the existing test_map_response_format to assert pass-through. Adds focused regression tests in tests/test_litellm/ covering preservation of type, strict, name, and schema body, plus that json_object alone still works. * fix(types): import Required from typing_extensions in gemini types * style: reformat sampling_handler.py for py312 black compat * refactor(mcp-sampling): extract helpers to fix PLR0915 too-many-statements in handle_sampling_create_message * fix(proxy-server): add explicit ProxyLogging type annotation to proxy_logging_obj to fix mypy inference * fix(mcp-sampling): suppress mypy assignment error on ImportError fallback for proxy_logging_obj * fix(test): use .value when comparing LlmProviders enum against string in test_default_api_base * fix(test): iterate LlmProviders enum in test_default_api_base to avoid str pollution from custom provider registration litellm.provider_list is a mutable global initialized to list(LlmProviders) but custom_llm_setup() appends plain provider strings to it. When a test_custom_llm.py test runs first in the same xdist worker, provider_list contains a str and calling .value on it raises AttributeError. Iterate the immutable LlmProviders enum instead, which is deterministic and what the check intends. * fix(mcp): depth-aware JSON-RPC response detection and neutral speed-priority fallback Replace the flat substring check in the truncated-body routing path with a top-level-key scan so a JSON-RPC response whose result payload nests a "method" field is still detected as a response and skips the session lock, removing a deadlock against the in-flight tool call awaiting it. Drop the inverse max_output_tokens speed proxy when no model exposes output_tokens_per_second; context-window size does not track latency, so a neutral score avoids biasing speedPriority toward the smallest-context model. * fix(guardrails): make ToolPermission rule reload atomic on invalid regex _load_rules appended each rule to self.rules before compiling its regex, so an invalid pattern raised mid-loop after the bad rule was already live but without a _compiled_rule_targets entry. _matches_regex reads a missing compiled target as a None pattern and returns True, turning the bad rule into a match-all that silently applies its decision to every tool. Via update_in_memory_litellm_params (PUT /guardrails) this corrupted the live guardrail. Build the parsed rules and compiled maps into locals and swap them in only after every regex compiles, and restore the previous ruleset if a live update is rejected, so an invalid regex now fails the update without leaving the guardrail enforcing a broken policy. * test(mcp): cover sampling conversion, model resolution, and elicitation relay paths The MCP sampling and elicitation handlers shipped with partial test coverage, leaving the response-to-MCP conversion, the model resolution fallback chain, completion-kwargs assembly, guardrail routing, and the entire elicitation relay untested. That pulled the PR's diff (patch) coverage below the codecov threshold even though overall project coverage rose. Add focused unit tests for _convert_openai_response_to_mcp_result, _convert_mcp_tools_to_openai, _convert_mcp_tool_choice_to_openai, image and audio content conversion, the hint-matching and fallback branches of _resolve_model_from_preferences, _build_completion_kwargs, the router and guardrail-rejection paths of _run_guardrails_and_call_llm, the handle_sampling_create_message success and error-propagation flows, the marker-hoisting fallback for tool content on unexpected roles, and the elicitation form/url/generic relay together with its decline paths --------- Co-authored-by: shin-berri Co-authored-by: yuneng-jiang Co-authored-by: lengkejun Co-authored-by: Yug Co-authored-by: Kent <72616338+kingdoooo@users.noreply.github.com> Co-authored-by: tanmay958 <53569547+tanmay958@users.noreply.github.com> Co-authored-by: DrishnaTrivedi <142084770+DrishnaTrivedi@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Navnit Shukla Co-authored-by: PRABHU KIRAN VANDRANKI <72809214+VANDRANKI@users.noreply.github.com> Co-authored-by: Adrian Lopez <109683617+adriangomez24@users.noreply.github.com> Co-authored-by: hcl Co-authored-by: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Co-authored-by: Dinesh Girbide <85330597+Dinesh-Girbide@users.noreply.github.com> Co-authored-by: cloudwiz <22098246+andrey-dubnik@users.noreply.github.com> Co-authored-by: Ahmad Khan Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + litellm/experimental_mcp_client/client.py | 115 +- .../get_supported_openai_params.py | 18 +- .../litellm_core_utils/streaming_handler.py | 26 + .../llms/azure/image_edit/transformation.py | 11 +- .../llms/bedrock_mantle/responses/__init__.py | 0 .../responses/transformation.py | 81 ++ .../llms/fireworks_ai/chat/transformation.py | 5 - litellm/llms/gemini/videos/transformation.py | 19 +- litellm/llms/huggingface/embedding/handler.py | 2 +- litellm/llms/snowflake/utils.py | 1 + .../vertex_ai_context_caching.py | 10 +- litellm/main.py | 10 +- ...odel_prices_and_context_window_backup.json | 38 + .../mcp_server/elicitation_handler.py | 163 +++ .../mcp_server/mcp_server_manager.py | 123 +- .../mcp_server/sampling_handler.py | 1279 +++++++++++++++++ .../proxy/_experimental/mcp_server/server.py | 672 ++++++--- litellm/proxy/caching_routes.py | 27 +- .../guardrail_hooks/tool_permission.py | 159 +- litellm/proxy/proxy_server.py | 14 +- litellm/types/llms/gemini.py | 12 +- .../types/mcp_server/mcp_server_manager.py | 2 + litellm/utils.py | 12 +- model_prices_and_context_window.json | 322 ++++- .../test_fireworks_ai_translation.py | 21 +- tests/local_testing/test_custom_llm.py | 88 +- tests/local_testing/test_get_llm_provider.py | 9 +- .../test_get_supported_openai_params.py | 134 ++ .../test_streaming_handler.py | 169 +++ .../test_azure_image_edit_transformation.py | 95 ++ ...bedrock_mantle_responses_transformation.py | 283 ++++ .../test_fireworks_ai_chat_transformation.py | 56 + .../test_gemini_video_transformation.py | 82 ++ .../test_huggingface_embedding_handler.py | 17 + .../test_vertex_ai_context_caching.py | 554 ++++++- .../test_mcp_elicitation_handler.py | 211 +++ .../mcp_server/test_mcp_hook_extra_headers.py | 31 +- .../test_mcp_sampling_completion_flow.py | 254 ++++ .../test_mcp_sampling_model_access.py | 327 +++++ .../test_mcp_sampling_model_resolution.py | 91 ++ .../test_mcp_sampling_priority_selection.py | 248 ++++ .../test_mcp_sampling_request_builder.py | 147 ++ .../test_mcp_sampling_response_conversion.py | 180 +++ .../test_mcp_sampling_tool_conversion.py | 312 ++++ .../mcp_server/test_mcp_server.py | 285 +++- .../mcp_server/test_mcp_server_manager.py | 24 +- .../mcp_server/test_mcp_stale_session.py | 34 +- .../guardrail_hooks/test_tool_permission.py | 153 +- .../test_litellm/proxy/test_caching_routes.py | 101 +- .../proxy/test_dynamic_mcp_route.py | 54 + tests/test_litellm/test_main.py | 18 +- tests/test_litellm/test_utils.py | 48 + 54 files changed, 6638 insertions(+), 517 deletions(-) create mode 100644 litellm/llms/bedrock_mantle/responses/__init__.py create mode 100644 litellm/llms/bedrock_mantle/responses/transformation.py create mode 100644 litellm/proxy/_experimental/mcp_server/elicitation_handler.py create mode 100644 litellm/proxy/_experimental/mcp_server/sampling_handler.py create mode 100644 tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py create mode 100644 tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 98c9dcb5ddf..e49f4a4699d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1740,6 +1740,9 @@ if TYPE_CHECKING: from .llms.openrouter.responses.transformation import ( OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig, ) + from .llms.bedrock_mantle.responses.transformation import ( + BedrockMantleResponsesAPIConfig as BedrockMantleResponsesAPIConfig, + ) from .llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index bdc3289b87c..5df8db7317d 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -237,6 +237,7 @@ LLM_CONFIG_NAMES = ( "PerplexityResponsesConfig", "DatabricksResponsesAPIConfig", "OpenRouterResponsesAPIConfig", + "BedrockMantleResponsesAPIConfig", "GoogleAIStudioInteractionsConfig", "OpenAIOSeriesConfig", "AnthropicSkillsConfig", @@ -958,6 +959,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.openrouter.responses.transformation", "OpenRouterResponsesAPIConfig", ), + "BedrockMantleResponsesAPIConfig": ( + ".llms.bedrock_mantle.responses.transformation", + "BedrockMantleResponsesAPIConfig", + ), "GoogleAIStudioInteractionsConfig": ( ".llms.gemini.interactions.transformation", "GoogleAIStudioInteractionsConfig", diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 7559fe142c4..0bc81ece5f0 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 +import os from typing import ( Any, Awaitable, @@ -16,7 +17,6 @@ from typing import ( TypeVar, Union, ) - import httpx from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client @@ -42,9 +42,8 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl - from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT +from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -67,7 +66,6 @@ TSessionResult = TypeVar("TSessionResult") class MCPSigV4Auth(httpx.Auth): """ httpx Auth class that signs each request with AWS SigV4. - This is used for MCP servers that require AWS SigV4 authentication, such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow() for every outgoing request, enabling per-request signature computation. @@ -92,10 +90,8 @@ class MCPSigV4Auth(httpx.Auth): "Missing botocore to use AWS SigV4 authentication. " "Run 'pip install boto3'." ) - self.service_name = aws_service_name or "bedrock-agentcore" self.region_name = aws_region_name or "us-east-1" - # Note: os.environ/ prefixed values are already resolved by # ProxyConfig._check_for_os_environ_vars() at config load time. # Values arrive here as plain strings. @@ -143,20 +139,17 @@ class MCPSigV4Auth(httpx.Auth): session_name = ( aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" ) - sts_kwargs: dict = {"region_name": aws_region_name} if aws_access_key_id and aws_secret_access_key: sts_kwargs["aws_access_key_id"] = aws_access_key_id sts_kwargs["aws_secret_access_key"] = aws_secret_access_key if aws_session_token: sts_kwargs["aws_session_token"] = aws_session_token - sts_client = boto3.client("sts", **sts_kwargs) sts_response = sts_client.assume_role( RoleArn=aws_role_name, RoleSessionName=session_name, ) - sts_creds = sts_response["Credentials"] return Credentials( access_key=sts_creds["AccessKeyId"], @@ -178,17 +171,14 @@ class MCPSigV4Auth(httpx.Auth): data=request.content, headers=dict(request.headers), ) - # Sign the request — SigV4Auth.add_auth() adds Authorization, # X-Amz-Date, and X-Amz-Security-Token (if session token present). # Host header is derived automatically from the URL. sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name) sigv4.add_auth(aws_request) - # Copy SigV4 headers back to the httpx request for header_name, header_value in aws_request.headers.items(): request.headers[header_name] = header_value - yield request @@ -198,6 +188,8 @@ class MCPClient: SSE and HTTP transports Authentication via Bearer token, Basic Auth, or API Key Tool calling with error handling and result parsing + Sampling callbacks for upstream server LLM requests + Elicitation callbacks for upstream server user-input requests """ def __init__( @@ -211,6 +203,9 @@ class MCPClient: extra_headers: Optional[Dict[str, str]] = None, ssl_verify: Optional[VerifyTypes] = None, aws_auth: Optional[httpx.Auth] = None, + sampling_callback: Optional[Callable] = None, + elicitation_callback: Optional[Callable] = None, + logging_callback: Optional[Callable] = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type @@ -222,6 +217,9 @@ class MCPClient: self.ssl_verify: Optional[VerifyTypes] = ssl_verify self._aws_auth: Optional[httpx.Auth] = aws_auth self._last_initialize_instructions: Optional[str] = None + self._sampling_callback: Optional[Callable] = sampling_callback + self._elicitation_callback: Optional[Callable] = elicitation_callback + self._logging_callback: Optional[Callable] = logging_callback # handle the basic auth value if provided if auth_value: self.update_auth_value(auth_value) @@ -231,23 +229,20 @@ class MCPClient: ) -> Tuple[Any, Optional[httpx.AsyncClient]]: """ Create the appropriate transport context based on transport type. - Returns: Tuple of (transport_context, http_client). http_client is only set for HTTP transport and needs cleanup. """ http_client: Optional[httpx.AsyncClient] = None - if self.transport_type == MCPTransport.stdio: if not self.stdio_config: raise ValueError("stdio_config is required for stdio transport") server_params = StdioServerParameters( command=self.stdio_config.get("command", ""), args=self.stdio_config.get("args", []), - env=self.stdio_config.get("env", {}), + env=self._get_safe_stdio_env(self.stdio_config.get("env")), ) return stdio_client(server_params), None - if self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() @@ -260,14 +255,12 @@ class MCPClient: ), None, ) - # HTTP transport (default) if streamable_http_client is None: raise ImportError( "streamable_http_client is not available. " "Please install mcp with HTTP support." ) - headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) @@ -281,6 +274,54 @@ class MCPClient: ) return transport_ctx, http_client + def _get_safe_stdio_env( + self, provided_env: Optional[Dict[str, str]] + ) -> Optional[Dict[str, str]]: + """ + Return a safe environment for the stdio subprocess. + + If provided_env is set, we use it as-is. + If provided_env is None, we return a minimal allowlist from the parent environment + to avoid leaking sensitive LiteLLM keys (OPENAI_API_KEY, etc.) to sub-processes. + """ + if provided_env is not None: + return provided_env + + # Minimal allowlist of safe/standard environment variables + safe_keys = { + "PATH", + "HOME", + "USER", + "LOGNAME", + "TMPDIR", + "TMP", + "TEMP", + "SHELL", + "LANG", + "LC_ALL", + # Node/Package manager caches + "NPM_CONFIG_CACHE", + "PNPM_HOME", + "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + # System info + "SYSTEMROOT", + "COMSPEC", + "PATHEXT", + "WINDIR", + } + + safe_env = {} + for key in safe_keys: + if key in os.environ: + safe_env[key] = os.environ[key] + + if "NPM_CONFIG_CACHE" not in safe_env: + safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR + + return safe_env + async def _execute_session_operation( self, transport_ctx: Any, @@ -288,13 +329,23 @@ class MCPClient: ) -> TSessionResult: """ Execute an operation within a transport and session context. - Handles entering/exiting contexts and running the operation. + Passes sampling/elicitation/logging callbacks to the ClientSession + so that upstream MCP servers can request LLM inference (sampling), + user input (elicitation), or send log messages. """ transport = await transport_ctx.__aenter__() try: read_stream, write_stream = transport[0], transport[1] - session_ctx = ClientSession(read_stream, write_stream) + # Build session kwargs with optional callbacks + session_kwargs: Dict[str, Any] = {} + if self._sampling_callback is not None: + session_kwargs["sampling_callback"] = self._sampling_callback + if self._elicitation_callback is not None: + session_kwargs["elicitation_callback"] = self._elicitation_callback + if self._logging_callback is not None: + session_kwargs["logging_callback"] = self._logging_callback + session_ctx = ClientSession(read_stream, write_stream, **session_kwargs) session = await session_ctx.__aenter__() try: init_result = await session.initialize() @@ -351,7 +402,6 @@ class MCPClient: def _get_auth_headers(self) -> dict: """Generate authentication headers based on auth type.""" headers = {} - if self._mcp_auth_value: if isinstance(self._mcp_auth_value, str): if self.auth_type == MCPAuth.bearer_token: @@ -373,17 +423,14 @@ class MCPClient: # Note: aws_sigv4 auth is not handled here — SigV4 requires per-request # signing (including the body hash), so it uses httpx.Auth flow instead # of static headers. See MCPSigV4Auth and _create_httpx_client_factory(). - # update the headers with the extra headers if self.extra_headers: headers.update(self.extra_headers) - return headers def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]: """ Create a custom httpx client factory that uses LiteLLM's SSL configuration. - This factory follows the same CA bundle path logic as http_handler.py: 1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle) 2. Check SSL_VERIFY environment variable @@ -400,17 +447,14 @@ class MCPClient: """Create an httpx.AsyncClient with LiteLLM's SSL configuration.""" # Get unified SSL configuration using the same logic as http_handler.py ssl_config = get_ssl_configuration(self.ssl_verify) - verbose_logger.debug( f"MCP client using SSL configuration: {type(ssl_config).__name__}" ) - # Use SigV4 auth if configured and no explicit auth provided. # The MCP SDK's sse_client and streamable_http_client call this # factory without passing auth=, so self._aws_auth is used. # For non-SigV4 clients, self._aws_auth is None — no behavior change. effective_auth = auth if auth is not None else self._aws_auth - return httpx.AsyncClient( headers=headers, timeout=timeout, @@ -458,7 +502,6 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( @@ -491,7 +534,6 @@ class MCPClient: f"MCP Tool '{call_tool_request_params.name}' progress: " f"{progress}/{total} ({percentage:.0f}%) - {message or ''}" ) - # Forward to Host if callback provided if host_progress_callback: try: @@ -521,7 +563,6 @@ class MCPClient: error_trace = traceback.format_exc() verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}") - # Log detailed error information error_type = type(e).__name__ verbose_logger.error( @@ -532,14 +573,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream - " "the MCP server may have crashed, disconnected, or timed out." ) - # Return a default error result instead of raising return MCPCallToolResult( content=[ @@ -577,14 +616,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_tools - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -617,7 +654,6 @@ class MCPClient: error_trace = traceback.format_exc() verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}") - # Log detailed error information error_type = type(e).__name__ verbose_logger.error( @@ -628,14 +664,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during get_prompt - " "the MCP server may have crashed, disconnected, or timed out." ) - raise async def list_resources(self) -> list[Resource]: @@ -667,14 +701,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_resources - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -709,14 +741,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_resource_templates - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -742,7 +772,6 @@ class MCPClient: error_trace = traceback.format_exc() verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}") - # Log detailed error information error_type = type(e).__name__ verbose_logger.error( @@ -753,12 +782,10 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during read_resource - " "the MCP server may have crashed, disconnected, or timed out." ) - raise diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index b8cdc8210fc..7c4f9941523 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -22,9 +22,11 @@ def get_supported_openai_params( # noqa: PLR0915 ``` Args: - base_model: For Azure, the true underlying model (e.g. ``"azure/gpt-5.2"``) - when the deployment name differs. Used for model-type detection so that - non-standard deployment names route to the correct config. + base_model: An optional capability hint for deployments whose ``model`` + label isn't recognized on its own (e.g. an Azure deployment name, or a + friendly Bedrock alias). It is additive: the result is the union of the + params supported by ``model`` and by ``base_model``, so a hint can only + add capabilities, never strip ones the real model already supports. Returns: - List if custom_llm_provider is mapped @@ -52,7 +54,15 @@ def get_supported_openai_params( # noqa: PLR0915 provider_config = None if provider_config and request_type == "chat_completion": - return provider_config.get_supported_openai_params(model=base_model or model) + supported_params = provider_config.get_supported_openai_params(model=model) + if base_model and base_model != model: + base_model_params = provider_config.get_supported_openai_params( + model=base_model + ) + supported_params = list( + dict.fromkeys([*supported_params, *base_model_params]) + ) + return supported_params if custom_llm_provider == "bedrock": return litellm.AmazonConverseConfig().get_supported_openai_params(model=model) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 55042a733ed..f3274151e5a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1149,6 +1149,32 @@ class CustomStreamWrapper: completion_obj: Dict[str, Any] = {"content": ""} from litellm.types.utils import GenericStreamingChunk as GChunk + if ( + isinstance(chunk, ModelResponseStream) + and self.custom_llm_provider is not None + and self.custom_llm_provider in litellm._custom_providers + ): + _has_content = bool( + chunk.choices + and chunk.choices[0].delta is not None + and ( + chunk.choices[0].delta.content + or chunk.choices[0].delta.tool_calls + ) + ) + if self.received_finish_reason is not None: + if not _has_content: + raise StopIteration + if chunk.choices and chunk.choices[0].finish_reason: + self.received_finish_reason = chunk.choices[0].finish_reason + if not _has_content: + return None + # Strip finish_reason from the content chunk so it appears + # only on the trailing empty-delta chunk (OpenAI spec). + # finish_reason_handler() will emit the proper terminal chunk. + chunk.choices[0].finish_reason = None # type: ignore[assignment] + return chunk + if ( isinstance(chunk, dict) and generic_chunk_has_all_required_fields( diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index a450ee0b217..72f1eef36c0 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -97,8 +97,15 @@ class AzureImageEditConfig(OpenAIImageEditConfig): ) original_url = httpx.URL(api_base) - # Extract api_version or use default - api_version = cast(Optional[str], litellm_params.get("api_version")) + # Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default. + # Mirrors the fallback chain used by the Azure chat path in common_utils.py, + # so callers that set a global / env api_version don't get an unversioned URL. + api_version = ( + cast(Optional[str], litellm_params.get("api_version")) + or litellm.api_version + or get_secret_str("AZURE_API_VERSION") + or litellm.AZURE_DEFAULT_API_VERSION + ) # Create a new dictionary with existing params query_params = dict(original_url.params) diff --git a/litellm/llms/bedrock_mantle/responses/__init__.py b/litellm/llms/bedrock_mantle/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py new file mode 100644 index 00000000000..b63fd0ecdb1 --- /dev/null +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -0,0 +1,81 @@ +""" +Amazon Bedrock Mantle - Responses API backend. + +gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses` +path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI +Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides +only the endpoint URL and Bearer authentication. + +Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the +standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4. +""" + +from typing import Optional + +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1" + +# Checked longest/most-specific first so a full endpoint URL collapses to host +# in one pass and the appended path never doubles. +_BASE_SUFFIXES_TO_STRIP = ( + "/openai/v1/responses", + "/v1/responses", + "/responses", + "/openai/v1", + "/v1", +) + + +class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.BEDROCK_MANTLE + + def get_complete_url( + self, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + region = ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + base = ( + api_base + or get_secret_str("BEDROCK_MANTLE_API_BASE") + or f"https://bedrock-mantle.{region}.api.aws" + ) + base = base.rstrip("/") + for suffix in _BASE_SUFFIXES_TO_STRIP: + if base.endswith(suffix): + base = base[: -len(suffix)] + break + return f"{base}/openai/v1/responses" + + def validate_environment( + self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] + ) -> dict: + litellm_params = litellm_params or GenericLiteLLMParams() + api_key = ( + litellm_params.api_key + or get_secret_str("BEDROCK_MANTLE_API_KEY") + or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + ) + if not api_key: + raise ValueError( + "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY " + "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key." + ) + headers["Authorization"] = f"Bearer {api_key}" + return headers + + def supports_native_file_search(self) -> bool: + return False + + def supports_native_websocket(self) -> bool: + return False diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 9e9d300b585..cca3b3da37a 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -170,11 +170,6 @@ class FireworksAIConfig(OpenAIGPTConfig): is_response_format_supported=False, enforce_tool_choice=False, # tools and response_format are both set, don't enforce tool_choice ) - elif "json_schema" in value: - optional_params["response_format"] = { - "type": "json_object", - "schema": value["json_schema"]["schema"], - } else: optional_params["response_format"] = value elif param == "max_completion_tokens": diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 77a95bfa5ab..644e96a7dd1 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -265,7 +265,11 @@ class GeminiVideoConfig(BaseVideoConfig): { "instances": [ { - "prompt": "A cat playing with a ball of yarn" + "prompt": "A cat playing with a ball of yarn", + "image": { + "bytesBase64Encoded": "...", + "mimeType": "image/jpeg" + } } ], "parameters": { @@ -275,13 +279,18 @@ class GeminiVideoConfig(BaseVideoConfig): } } """ - instance = GeminiVideoGenerationInstance(prompt=prompt) + instance: GeminiVideoGenerationInstance = {"prompt": prompt} params_copy = video_create_optional_request_params.copy() - if "image" in params_copy and params_copy["image"] is not None: - image_data = _convert_image_to_gemini_format(params_copy["image"]) - params_copy["image"] = image_data + if "image" in params_copy: + image = params_copy.pop("image") + if image is not None: + if isinstance(image, dict): + image_data = image + else: + image_data = _convert_image_to_gemini_format(image) + instance["image"] = image_data parameters = GeminiVideoGenerationParameters(**params_copy) diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 226f6b2ebad..6be885b1f91 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -239,7 +239,7 @@ class HuggingFaceEmbedding(BaseLLM): model_response.model = model input_tokens = 0 for text in input: - input_tokens += len(encoding.encode(text)) + input_tokens += len(encoding.encode(text, disallowed_special=())) setattr( model_response, diff --git a/litellm/llms/snowflake/utils.py b/litellm/llms/snowflake/utils.py index d84efdd9fcd..4f79006f6f8 100644 --- a/litellm/llms/snowflake/utils.py +++ b/litellm/llms/snowflake/utils.py @@ -25,6 +25,7 @@ class SnowflakeBaseConfig: "temperature", "max_tokens", "top_p", + "stream", "response_format", "tools", "tool_choice", diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 3f945adca0d..e9f08f403f9 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -337,6 +337,7 @@ class ContextCachingEndpoints(VertexBase): return messages, optional_params, None tools = optional_params.pop("tools", None) + tool_choice = optional_params.pop("tool_choice", None) ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( @@ -371,7 +372,7 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, model=model + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) google_cache_name = self.check_cache( cache_key=generated_cache_key, @@ -402,6 +403,8 @@ class ContextCachingEndpoints(VertexBase): ) cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice ## LOGGING logging_obj.pre_call( @@ -487,6 +490,7 @@ class ContextCachingEndpoints(VertexBase): return messages, optional_params, None tools = optional_params.pop("tools", None) + tool_choice = optional_params.pop("tool_choice", None) ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( @@ -518,7 +522,7 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, model=model + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) google_cache_name = await self.async_check_cache( cache_key=generated_cache_key, @@ -550,6 +554,8 @@ class ContextCachingEndpoints(VertexBase): ) cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice ## LOGGING logging_obj.pre_call( diff --git a/litellm/main.py b/litellm/main.py index 96f81381c86..c8aae0ce85b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1322,7 +1322,9 @@ def completion( # type: ignore # noqa: PLR0915 preset_cache_key = kwargs.get("preset_cache_key", None) hf_model_name = kwargs.get("hf_model_name", None) supports_system_message = kwargs.get("supports_system_message", None) - base_model = kwargs.get("base_model", None) + base_model = kwargs.get("base_model", None) or ( + model_info.get("base_model") if isinstance(model_info, dict) else None + ) ### DISABLE FLAGS ### disable_add_transform_inline_image_block = kwargs.get( "disable_add_transform_inline_image_block", None @@ -1534,11 +1536,7 @@ def completion( # type: ignore # noqa: PLR0915 "logit_bias": logit_bias, "user": user, # params to identify the model - "model": ( - model_info.get("base_model") - if isinstance(model_info, dict) and model_info.get("base_model") - else model - ), + "model": model, "custom_llm_provider": custom_llm_provider, "response_format": response_format, "seed": seed, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ed6de4fa6b7..fbe6c097202 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41223,6 +41223,44 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock_mantle/openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "cache_read_input_token_cost": 5.5e-07, + "output_cost_per_token": 3.3e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": ["/v1/responses"], + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["text"], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "cache_read_input_token_cost": 2.75e-07, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": ["/v1/responses"], + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["text"], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py new file mode 100644 index 00000000000..e42270bf10b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py @@ -0,0 +1,163 @@ +""" +MCP Elicitation Handler +Handles `elicitation/create` requests from upstream MCP servers by either: +1. Relaying them to the connected downstream MCP client (if it supports elicitation) +2. Returning a decline/error response (if no downstream client or unsupported) +Supports both Form mode (structured data collection) and URL mode (external URL +navigation for sensitive interactions like OAuth). +MCP Spec Reference: + https://modelcontextprotocol.io/specification/2025-11-25/client/elicitation +""" + +from typing import Any, Optional, Union +from litellm._logging import verbose_logger + +# Guard imports that require the mcp package +try: + from mcp.types import ( + ElicitRequestFormParams, + ElicitRequestParams, + ElicitRequestURLParams, + ElicitResult, + ErrorData, + ) + + MCP_ELICITATION_AVAILABLE = True +except ImportError: + MCP_ELICITATION_AVAILABLE = False + + +async def handle_elicitation_request( + context: Any, + params: "ElicitRequestParams", + downstream_session: Optional[Any] = None, + downstream_capabilities: Optional[Any] = None, +) -> Union["ElicitResult", "ErrorData"]: + """ + Handle an MCP elicitation/create request from an upstream MCP server. + In Gateway mode (Mode A), we relay the elicitation request to the + connected downstream client if they declared elicitation capabilities. + In Tool Bridge mode (Mode B), there's no persistent downstream MCP + client, so we return a decline response. + Args: + context: MCP RequestContext from the upstream server connection. + params: The ElicitRequestParams (either form or URL mode). + downstream_session: The ServerSession to the downstream client, + if available (for relaying). + downstream_capabilities: The downstream client's declared + capabilities, used to check elicitation support. + Returns: + ElicitResult with the user's response, or ErrorData on failure. + """ + if not MCP_ELICITATION_AVAILABLE: + return ErrorData( + code=-1, + message="MCP elicitation is not available (mcp package not installed)", + ) + try: + mode = getattr(params, "mode", "form") + verbose_logger.info( + "MCP elicitation: received request mode=%s, message=%s", + mode, + getattr(params, "message", ""), + ) + # Check if we have a downstream session to relay to + if downstream_session is not None: + return await _relay_elicitation_to_downstream( + params=params, + downstream_session=downstream_session, + downstream_capabilities=downstream_capabilities, + ) + # No downstream session — we're in Tool Bridge mode + # or the client doesn't support elicitation + verbose_logger.info( + "MCP elicitation: no downstream session available, declining" + ) + return ElicitResult( + action="decline", + ) + except Exception as e: + verbose_logger.exception("MCP elicitation handler failed: %s", e) + return ErrorData( + code=-1, + message=f"Elicitation failed: {str(e)}", + ) + + +async def _relay_elicitation_to_downstream( + params: "ElicitRequestParams", + downstream_session: Any, + downstream_capabilities: Optional[Any] = None, +) -> Union["ElicitResult", "ErrorData"]: + """ + Relay an elicitation request to the downstream MCP client. + Uses the ServerSession's elicit_form() or elicit_url() methods to + send the elicitation request back to the connected client. + Args: + params: The elicitation request parameters. + downstream_session: The ServerSession connected to the downstream client. + downstream_capabilities: Client capabilities to check support. + Returns: + ElicitResult from the downstream client. + """ + mode = getattr(params, "mode", "form") + # Check if the downstream client supports the requested mode + if downstream_capabilities is not None: + elicit_caps = getattr(downstream_capabilities, "elicitation", None) + if elicit_caps is None: + verbose_logger.info( + "MCP elicitation: downstream client does not support elicitation" + ) + return ElicitResult(action="decline") + if mode == "url": + url_cap = getattr(elicit_caps, "url", None) + if url_cap is None: + verbose_logger.info( + "MCP elicitation: downstream client does not support URL mode" + ) + return ElicitResult(action="decline") + if mode == "form": + form_cap = getattr(elicit_caps, "form", None) + if form_cap is None: + verbose_logger.info( + "MCP elicitation: downstream client does not support form mode" + ) + return ElicitResult(action="decline") + try: + if mode == "url" and isinstance(params, ElicitRequestURLParams): + # URL mode: relay URL to client for external navigation + verbose_logger.info( + "MCP elicitation: relaying URL mode to downstream, url=%s", + getattr(params, "url", ""), + ) + result = await downstream_session.elicit_url( + message=params.message, + url=params.url, + elicitation_id=getattr(params, "elicitationId", None), + ) + elif isinstance(params, ElicitRequestFormParams): + # Form mode: relay structured form to client + verbose_logger.info("MCP elicitation: relaying form mode to downstream") + result = await downstream_session.elicit_form( + message=params.message, + requestedSchema=getattr(params, "requestedSchema", None), + ) + else: + # Fallback for generic ElicitRequestParams — pass an empty schema + # since elicit() requires requestedSchema as a positional arg. + verbose_logger.info( + "MCP elicitation: relaying generic elicitation to downstream" + ) + result = await downstream_session.elicit( + message=getattr(params, "message", ""), + requestedSchema=getattr(params, "requestedSchema", {}), + ) + verbose_logger.info( + "MCP elicitation: downstream responded with action=%s", + getattr(result, "action", "unknown"), + ) + return result + except Exception as e: + verbose_logger.warning("MCP elicitation: failed to relay to downstream: %s", e) + # If relay fails, decline gracefully + return ElicitResult(action="decline") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d9b112f6c21..0d2008cdade 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -49,6 +49,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + MCP_ELICITATION_AVAILABLE, +) +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + MCP_SAMPLING_AVAILABLE, +) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, @@ -289,6 +295,82 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: return data +def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): + """ + Create a sampling callback for MCP ClientSession. + Returns a callable that handles sampling/createMessage requests from + upstream MCP servers by routing them through litellm.acompletion(). + """ + if not MCP_SAMPLING_AVAILABLE: + return None + + async def _sampling_callback(context, params): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + import litellm + from litellm.proxy._experimental.mcp_server.server import ( + get_active_auth_context, + ) + + auth_context = get_active_auth_context() + resolved_auth = user_api_key_auth or ( + auth_context.user_api_key_auth if auth_context else None + ) + # Forward original HTTP headers and client IP so that + # header-dependent guardrails, tag-based routing, trace + # correlation, and forward_llm_provider_auth_headers work + # correctly for sampling sub-calls. + _raw_headers = getattr(auth_context, "raw_headers", None) + _client_ip = getattr(auth_context, "client_ip", None) + + return await handle_sampling_create_message( + context=context, + params=params, + default_model=getattr(litellm, "default_mcp_sampling_model", None), + user_api_key_auth=resolved_auth, + raw_headers=_raw_headers, + client_ip=_client_ip, + ) + + return _sampling_callback + + +def _create_elicitation_callback(): + """ + Create an elicitation callback for MCP ClientSession. + Returns a callable that handles elicitation/create requests from + upstream MCP servers. In gateway mode, this relays to the downstream + client; in tool bridge mode, it returns a decline response. + """ + if not MCP_ELICITATION_AVAILABLE: + return None + + async def _elicitation_callback(context, params): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session + + # In Gateway mode, we relay the elicitation request to the downstream client + # that triggered the current operation. + downstream_session = get_active_mcp_session() + downstream_capabilities = ( + getattr(downstream_session, "capabilities", None) + if downstream_session + else None + ) + + return await handle_elicitation_request( + context=context, + params=params, + downstream_session=downstream_session, + downstream_capabilities=downstream_capabilities, + ) + + return _elicitation_callback + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -600,6 +682,8 @@ class MCPServerManager: "subject_token_type", "urn:ietf:params:oauth:token-type:access_token", ), + allow_sampling=bool(server_config.get("allow_sampling", False)), + allow_elicitation=bool(server_config.get("allow_elicitation", False)), ) self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") @@ -699,8 +783,7 @@ class MCPServerManager: ) verbose_logger.debug( - f"Using headers for OpenAPI tools (excluding sensitive values): " - f"{list(headers.keys())}" + f"Using headers for OpenAPI tools (excluding sensitive values): {list(headers.keys())}" ) # Extract and register tools from OpenAPI paths @@ -1494,6 +1577,7 @@ class MCPServerManager: extra_headers: Optional[Dict[str, str]] = None, stdio_env: Optional[Dict[str, str]] = None, subject_token: Optional[str] = None, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> MCPClient: """ Create an MCPClient instance for the given server. @@ -1510,6 +1594,7 @@ class MCPServerManager: extra_headers: Additional headers to forward. stdio_env: Environment variables for stdio transport. subject_token: Optional user JWT for token exchange (OBO) flow. + user_api_key_auth: Optional auth context for sampling callbacks. Returns: Configured MCP client instance. @@ -1520,23 +1605,44 @@ class MCPServerManager: transport = server.transport or MCPTransport.sse + # Create sampling and elicitation callbacks for this client + sampling_cb = ( + _create_sampling_callback(user_api_key_auth=user_api_key_auth) + if server.allow_sampling + else None + ) + elicitation_cb = ( + _create_elicitation_callback() if server.allow_elicitation else None + ) + # Handle stdio transport if transport == MCPTransport.stdio: resolved_env = ( - stdio_env if stdio_env is not None else dict(server.env or {}) + stdio_env + if stdio_env is not None + else (dict(server.env) if server.env is not None else None) ) # Ensure npm-based STDIO MCP servers have a writable cache dir. # In containers the default (~/.npm or /app/.npm) may not exist # or be read-only, causing npx to fail with ENOENT. - if "NPM_CONFIG_CACHE" not in resolved_env: + if resolved_env is not None and "NPM_CONFIG_CACHE" not in resolved_env: resolved_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR # Defense-in-depth: block commands not in the allowlist. # The Pydantic validator blocks new servers; this catches legacy # config/DB records predating the allowlist. if server.command: base_command = os.path.basename(server.command) - if base_command not in MCP_STDIO_ALLOWED_COMMANDS: + # Strip .exe/.cmd/.bat/.com suffix for Windows compatibility + base_command_no_ext = base_command.lower() + for ext in [".exe", ".cmd", ".bat", ".com"]: + if base_command.lower().endswith(ext): + base_command_no_ext = base_command[: -len(ext)].lower() + break + if ( + base_command.lower() not in MCP_STDIO_ALLOWED_COMMANDS + and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS + ): raise HTTPException( status_code=403, detail=f"MCP stdio command '{server.command}' is not in the allowlist ({sorted(MCP_STDIO_ALLOWED_COMMANDS)}). " @@ -1559,6 +1665,8 @@ class MCPServerManager: timeout=MCP_CLIENT_TIMEOUT, stdio_config=stdio_config, extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) else: # For HTTP/SSE transports @@ -1585,6 +1693,8 @@ class MCPServerManager: timeout=MCP_CLIENT_TIMEOUT, extra_headers=extra_headers, aws_auth=aws_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) async def _get_tools_from_server( @@ -1668,6 +1778,7 @@ class MCPServerManager: mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, stdio_env=stdio_env, + user_api_key_auth=user_api_key_auth, ) ## HANDLE OPENAPI TOOLS @@ -3030,6 +3141,7 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, ) call_tool_params = MCPCallToolRequestParams( @@ -3260,7 +3372,6 @@ class MCPServerManager: ) ) else: - # For regular MCP servers, use the MCP client return await self._call_regular_mcp_tool( mcp_server=mcp_server, original_tool_name=name, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py new file mode 100644 index 00000000000..1637c9eb0b9 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -0,0 +1,1279 @@ +""" +MCP Sampling Handler +Handles `sampling/createMessage` requests from upstream MCP servers by +routing them through LiteLLM's internal completion infrastructure. +This allows MCP servers to perform agentic reasoning (e.g., multi-step +tool calling, chain-of-thought) without needing their own LLM API keys — +LiteLLM acts as the LLM provider using its existing 100+ provider support, +cost tracking, rate limiting, and model routing. +MCP Spec Reference: + https://modelcontextprotocol.io/specification/2025-11-25/client/sampling +""" + +from typing import Any, Dict, List, Optional, Union +import typing + +if typing.TYPE_CHECKING: + from litellm.proxy.utils import ProxyLogging + +from litellm._logging import verbose_logger + +from fastapi import HTTPException + +# Guard imports that require the mcp package +try: + from mcp.types import ( + CreateMessageRequestParams, + CreateMessageResult, + CreateMessageResultWithTools, + ErrorData, + ModelPreferences, + SamplingMessage, + TextContent, + Tool, + ToolChoice, + ToolUseContent, + ) + + MCP_SAMPLING_AVAILABLE = True +except ImportError as _sampling_import_err: + MCP_SAMPLING_AVAILABLE = False + verbose_logger.warning( + "MCP sampling disabled: failed to import required types from mcp.types — %s. " + "This usually means the 'mcp' package is not installed or is an older version " + "that does not support sampling. Install/upgrade with: pip install 'mcp>=1.1'", + _sampling_import_err, + ) + + +def _resolve_model_from_preferences( + model_preferences: Optional["ModelPreferences"], + default_model: Optional[str] = None, +) -> str: + """ + Resolve an LLM model name from MCP ModelPreferences. + Strategy: + 1. Check hints for substring matches against known model names. + 2. Fall back to priority-based selection (cost/speed/intelligence). + 3. Fall back to the configured default model. + Args: + model_preferences: MCP ModelPreferences with hints and priorities. + default_model: Fallback model if no hint matches. + Returns: + A model string suitable for litellm.acompletion(). + """ + import litellm + + # Build list of available model names from proxy Router or litellm.model_list + available_model_names: list = [] + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is not None: + available_model_names = llm_router.get_model_names() + except Exception: + pass + if not available_model_names and litellm.model_list: + for entry in litellm.model_list: + if isinstance(entry, dict): + name = entry.get("model_name") + if name: + available_model_names.append(name) + elif isinstance(entry, str): + available_model_names.append(entry) + if model_preferences and model_preferences.hints: + for hint in model_preferences.hints: + hint_name = getattr(hint, "name", None) + if not hint_name: + continue + # Try direct match first + if hint_name in available_model_names: + verbose_logger.debug( + "MCP sampling model resolution: direct hint match '%s'", + hint_name, + ) + return hint_name + # Try substring match against known models + for model_name in available_model_names: + if hint_name.lower() in model_name.lower(): + verbose_logger.debug( + "MCP sampling model resolution: substring hint match " + "'%s' -> '%s'", + hint_name, + model_name, + ) + return model_name + verbose_logger.debug( + "MCP sampling model resolution: no hint matched from %s " + "against %d available models", + [getattr(h, "name", None) for h in model_preferences.hints], + len(available_model_names), + ) + + # 2. Priority-based selection (cost/speed/intelligence) + if ( + model_preferences + and available_model_names + and _has_priorities(model_preferences) + ): + best = _select_model_by_priority(available_model_names, model_preferences) + if best is not None: + verbose_logger.debug( + "MCP sampling model resolution: priority-based selection chose '%s'", + best, + ) + return best + + # 3. Use default model from caller + if default_model: + verbose_logger.debug( + "MCP sampling model resolution: using caller-provided default '%s'", + default_model, + ) + return default_model + # Fall back to first available model + if available_model_names: + verbose_logger.debug( + "MCP sampling model resolution: no default configured, " + "falling back to first available model '%s'", + available_model_names[0], + ) + return available_model_names[0] + # Last resort - use LiteLLM default or raise error + default_sampling_model = getattr(litellm, "default_mcp_sampling_model", None) + if default_sampling_model: + verbose_logger.debug( + "MCP sampling model resolution: using litellm.default_mcp_sampling_model='%s'", + default_sampling_model, + ) + return default_sampling_model + raise ValueError( + "No model could be resolved for MCP sampling. Please configure 'default_mcp_sampling_model' in your LiteLLM configuration." + ) + + +def _has_priorities(model_preferences: "ModelPreferences") -> bool: + """Return True if any priority weight is set (non-None and > 0).""" + return any( + (getattr(model_preferences, attr, None) or 0) > 0 + for attr in ("costPriority", "speedPriority", "intelligencePriority") + ) + + +def _select_model_by_priority( + model_names: List[str], + model_preferences: "ModelPreferences", +) -> Optional[str]: + """Score available models by MCP priority weights and return the best. + + Scoring strategy (per the MCP spec, priorities are 0-1 floats): + + * **costPriority** — higher means "prefer cheaper models". + Metric: combined (input + output) cost per token from + ``model_prices_and_context_window.json``. Lower cost → higher score. + + * **speedPriority** — higher means "prefer faster models". + Metric: ``output_tokens_per_second`` from model info when available; + otherwise a neutral score for every candidate, since no reliable + latency proxy exists (context-window size does not track speed). + + * **intelligencePriority** — higher means "prefer smarter models". + Metric: ``max_output_tokens`` is used as a rough capability proxy + (frontier models expose larger context windows). + + Each metric is min-max normalised across the candidate set so that + every model gets a 0-1 score per dimension. The final score is the + weighted sum of the three normalised dimensions. + + Returns the highest-scoring model name, or None if scoring fails for + all candidates (e.g. no model_info available). + """ + import litellm as _litellm + + cost_weight = getattr(model_preferences, "costPriority", None) or 0.0 + speed_weight = getattr(model_preferences, "speedPriority", None) or 0.0 + intel_weight = getattr(model_preferences, "intelligencePriority", None) or 0.0 + + # Gather raw metrics for each model + scored: List[Dict[str, Any]] = [] + for name in model_names: + try: + info = _litellm.get_model_info(name) + except Exception: + continue + input_cost = info.get("input_cost_per_token") or 0.0 + output_cost = info.get("output_cost_per_token") or 0.0 + total_cost = input_cost + output_cost + max_output = info.get("max_output_tokens") or info.get("max_tokens") or 0 + output_tps = info.get("output_tokens_per_second") or 0.0 + scored.append( + { + "name": name, + "cost": total_cost, + "max_output": max_output, + "output_tps": output_tps, + } + ) + + if not scored: + return None + + # Min-max normalisation helpers + def _normalise(values: List[float], invert: bool = False) -> List[float]: + """Normalise to [0, 1]. If *invert*, lower raw → higher score.""" + lo, hi = min(values), max(values) + if hi == lo: + return [0.5] * len(values) # all equal → neutral score + normed = [(v - lo) / (hi - lo) for v in values] + if invert: + normed = [1.0 - n for n in normed] + return normed + + costs = [s["cost"] for s in scored] + max_outputs = [float(s["max_output"]) for s in scored] + output_tps_values = [s["output_tps"] for s in scored] + + # costPriority: lower cost → higher score (invert) + cost_scores = _normalise(costs, invert=True) + # speedPriority: use output_tokens_per_second if any model has it, + # otherwise a neutral score (no reliable latency proxy is available). + if any(v > 0 for v in output_tps_values): + speed_scores = _normalise(output_tps_values, invert=False) + else: + speed_scores = [0.5] * len(scored) + # intelligencePriority: higher max_output → smarter + intel_scores = _normalise(max_outputs, invert=False) + + best_name = None + best_score = -1.0 + for i, entry in enumerate(scored): + score = ( + cost_weight * cost_scores[i] + + speed_weight * speed_scores[i] + + intel_weight * intel_scores[i] + ) + verbose_logger.debug( + "MCP priority scoring: model=%s cost_score=%.3f speed_score=%.3f " + "intel_score=%.3f → weighted=%.3f", + entry["name"], + cost_scores[i], + speed_scores[i], + intel_scores[i], + score, + ) + if score > best_score: + best_score = score + best_name = entry["name"] + + return best_name + + +def _convert_mcp_content_to_openai( + content: Any, +) -> Union[str, Dict[str, Any], List[Dict[str, Any]]]: + """ + Convert MCP SamplingMessage content to OpenAI message content format. + Handles: + - TextContent → string or {"type": "text", "text": ...} + - ImageContent → {"type": "image_url", "image_url": {"url": "data:..."}} + - AudioContent → {"type": "input_audio", "input_audio": {...}} + - ToolUseContent → function call representation + - ToolResultContent → tool result representation + - List of mixed content → list of content parts + """ + if isinstance(content, list): + parts = [] + for item in content: + converted = _convert_single_content(item) + if isinstance(converted, list): + parts.extend(converted) + else: + parts.append(converted) + return parts + return _convert_single_content(content) + + +def _convert_single_content( + content: Any, +) -> Union[Dict[str, Any], List[Dict[str, Any]]]: + """Convert a single MCP content item to OpenAI format. + + For text/image/audio content, returns a single content-part dict. + For tool_use/tool_result, returns a dict with a ``_marker_type`` key + so the caller (``_convert_mcp_messages_to_openai``) can hoist it to + the correct message-level position (``tool_calls`` array or a + separate ``role: "tool"`` message). + """ + import json + + content_type = getattr(content, "type", None) + if content_type == "text": + return {"type": "text", "text": content.text} + elif content_type == "image": + data = getattr(content, "data", "") + mime_type = getattr(content, "mimeType", "image/png") + return { + "type": "image_url", + "image_url": {"url": f"data:{mime_type};base64,{data}"}, + } + elif content_type == "audio": + data = getattr(content, "data", "") + mime_type = getattr(content, "mimeType", "audio/wav") + # Map MIME type to OpenAI audio format + format_map = { + "audio/wav": "wav", + "audio/mp3": "mp3", + "audio/mpeg": "mp3", + "audio/flac": "flac", + "audio/ogg": "ogg", + } + audio_format = format_map.get(mime_type, "wav") + return { + "type": "input_audio", + "input_audio": {"data": data, "format": audio_format}, + } + elif content_type == "tool_use": + # ToolUseContent → proper OpenAI function-call representation. + # The ``_marker_type`` key lets the message-level converter + # hoist this into the ``tool_calls`` array on the assistant + # message instead of embedding it inline as a content part. + return { + "_marker_type": "tool_use", + "id": getattr(content, "id", f"call_{id(content)}"), + "type": "function", + "function": { + "name": getattr(content, "name", ""), + "arguments": json.dumps(getattr(content, "input", {}), default=str), + }, + } + elif content_type == "tool_result": + # ToolResultContent → proper OpenAI tool-role message. + # Marked so the message-level converter can emit it as a + # separate ``{"role": "tool", ...}`` message. + tool_use_id = getattr(content, "toolUseId", "") + nested_content = getattr(content, "content", []) + if isinstance(nested_content, list): + text_parts = [ + getattr(c, "text", str(c)) + for c in nested_content + if getattr(c, "type", None) == "text" + ] + result_text = "\n".join(text_parts) if text_parts else "" + else: + result_text = str(nested_content) + return { + "_marker_type": "tool_result", + "role": "tool", + "tool_call_id": tool_use_id, + "content": result_text, + } + # Fallback: treat as text + return {"type": "text", "text": str(content)} + + +def _convert_mcp_messages_to_openai( + messages: List["SamplingMessage"], + system_prompt: Optional[str] = None, +) -> List[Dict[str, Any]]: + """ + Convert MCP SamplingMessage list to OpenAI messages format. + MCP messages use: + - role: "user" | "assistant" + - content: TextContent | ImageContent | AudioContent | ToolUseContent + | ToolResultContent | list[...] + OpenAI messages use: + - role: "system" | "user" | "assistant" | "tool" + - content: str | list[content_part] + """ + openai_messages: List[Dict[str, Any]] = [] + # Add system prompt if provided + if system_prompt: + openai_messages.append({"role": "system", "content": system_prompt}) + for msg in messages: + role = msg.role + content = msg.content + # Handle tool use content from assistant + if role == "assistant" and _has_tool_use(content): + tool_calls = _extract_tool_calls(content) + if tool_calls: + openai_msg: Dict[str, Any] = { + "role": "assistant", + "tool_calls": tool_calls, + } + # Also include any text content alongside tool calls + text_parts = _extract_text_parts(content) + if text_parts: + openai_msg["content"] = text_parts + openai_messages.append(openai_msg) + continue + # Handle tool result content from user + if role == "user" and _has_tool_result(content): + tool_results = _extract_tool_results(content) + for tool_result in tool_results: + openai_messages.append(tool_result) + continue + # Standard text/image/audio message — also handles any stray + # tool_use / tool_result that slipped past the fast-path checks + # above (e.g. unexpected role, single non-list content). + converted = _convert_mcp_content_to_openai(content) + converted_parts = ( + converted + if isinstance(converted, list) + else ([converted] if isinstance(converted, dict) else []) + ) + + # Separate marker items from regular content parts + tool_call_markers = [] + tool_result_markers = [] + regular_parts = [] + for part in converted_parts: + marker = part.get("_marker_type") if isinstance(part, dict) else None + if marker == "tool_use": + # Strip the internal marker before emitting + tc = {k: v for k, v in part.items() if k != "_marker_type"} + tool_call_markers.append(tc) + elif marker == "tool_result": + tr = {k: v for k, v in part.items() if k != "_marker_type"} + tool_result_markers.append(tr) + else: + regular_parts.append(part) + + # Emit assistant message with tool_calls if any were found + if tool_call_markers: + openai_msg_tc: Dict[str, Any] = { + "role": "assistant", + "tool_calls": tool_call_markers, + } + if regular_parts: + openai_msg_tc["content"] = regular_parts + openai_messages.append(openai_msg_tc) + elif regular_parts: + if isinstance(converted, str): + openai_messages.append({"role": role, "content": converted}) + else: + openai_messages.append({"role": role, "content": regular_parts}) + + # Emit separate tool-result messages + for tr in tool_result_markers: + openai_messages.append(tr) + + return openai_messages + + +def _has_tool_use(content: Any) -> bool: + """Check if content contains ToolUseContent.""" + if isinstance(content, list): + return any(getattr(c, "type", None) == "tool_use" for c in content) + return getattr(content, "type", None) == "tool_use" + + +def _has_tool_result(content: Any) -> bool: + """Check if content contains ToolResultContent.""" + if isinstance(content, list): + return any(getattr(c, "type", None) == "tool_result" for c in content) + return getattr(content, "type", None) == "tool_result" + + +def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]: + """Extract OpenAI-format tool_calls from MCP ToolUseContent.""" + import json + + items = content if isinstance(content, list) else [content] + tool_calls = [] + for item in items: + if getattr(item, "type", None) == "tool_use": + tool_calls.append( + { + "id": getattr(item, "id", f"call_{id(item)}"), + "type": "function", + "function": { + "name": getattr(item, "name", ""), + "arguments": json.dumps( + getattr(item, "input", {}), default=str + ), + }, + } + ) + return tool_calls + + +def _extract_text_parts(content: Any) -> Optional[str]: + """Extract text parts from mixed content.""" + items = content if isinstance(content, list) else [content] + texts = [] + for item in items: + if getattr(item, "type", None) == "text": + texts.append(getattr(item, "text", "")) + return "\n".join(texts) if texts else None + + +def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: + """Extract OpenAI-format tool messages from MCP ToolResultContent.""" + items = content if isinstance(content, list) else [content] + results = [] + for item in items: + if getattr(item, "type", None) == "tool_result": + tool_use_id = getattr(item, "toolUseId", "") + # Extract text from nested content + nested_content = getattr(item, "content", []) + if isinstance(nested_content, list): + text_parts = [ + getattr(c, "text", str(c)) + for c in nested_content + if getattr(c, "type", None) == "text" + ] + result_text = "\n".join(text_parts) if text_parts else "" + else: + result_text = str(nested_content) + results.append( + { + "role": "tool", + "tool_call_id": tool_use_id, + "content": result_text, + } + ) + return results + + +def _convert_mcp_tools_to_openai( + tools: Optional[List["Tool"]], +) -> Optional[List[Dict[str, Any]]]: + """ + Convert MCP Tool definitions to OpenAI function calling format. + MCP Tool: {name, description, inputSchema} + OpenAI Tool: {type: "function", function: {name, description, parameters}} + """ + if not tools: + return None + openai_tools = [] + for tool in tools: + openai_tool = { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description or "", + "parameters": tool.inputSchema + or { + "type": "object", + "properties": {}, + }, + }, + } + openai_tools.append(openai_tool) + return openai_tools + + +def _convert_mcp_tool_choice_to_openai( + tool_choice: Optional["ToolChoice"], +) -> Optional[Union[str, Dict[str, Any]]]: + """ + Convert MCP ToolChoice to OpenAI tool_choice format. + MCP: {mode: "auto"} | {mode: "required"} | {mode: "none"} + OpenAI: "auto" | "required" | "none" + """ + if not tool_choice: + return None + mode = getattr(tool_choice, "mode", "auto") + if mode == "auto": + return "auto" + elif mode == "required": + return "required" + elif mode == "none": + return "none" + return "auto" + + +def _convert_openai_response_to_mcp_result( + response: Any, + model_name: str, +) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: + """ + Convert a litellm completion response to MCP CreateMessageResult. + Args: + response: The litellm ModelResponse. + model_name: The model that was used. + Returns: + MCP CreateMessageResult or CreateMessageResultWithTools. + """ + if not response.choices: + verbose_logger.warning( + "MCP sampling: LLM returned empty choices list for model=%s " + "(possible content filter or provider error)", + model_name, + ) + return ErrorData( + code=-1, + message=( + f"LLM returned no choices for model '{model_name}'. " + "This may indicate content filtering or a provider-side error." + ), + ) + choice = response.choices[0] + message = choice.message + # Determine stop reason + finish_reason = getattr(choice, "finish_reason", "stop") + if finish_reason == "tool_calls": + stop_reason = "toolUse" + elif finish_reason == "length": + stop_reason = "maxTokens" + else: + stop_reason = "endTurn" + actual_model = getattr(response, "model", model_name) or model_name + # Check if response has tool calls + tool_calls = getattr(message, "tool_calls", None) + if tool_calls: + # Build ToolUseContent items + content_parts: "List[Any]" = [] + # Include text content if present + if message.content: + content_parts.append(TextContent(type="text", text=message.content)) + # Convert tool calls to MCP ToolUseContent + for tc in tool_calls: + import json + + tool_input = tc.function.arguments + if isinstance(tool_input, str): + try: + tool_input = json.loads(tool_input) + except (json.JSONDecodeError, TypeError): + tool_input = {"raw": tool_input} + content_parts.append( + ToolUseContent( + type="tool_use", + id=tc.id, + name=tc.function.name, + input=tool_input, + ) + ) + return CreateMessageResultWithTools( + role="assistant", + content=content_parts, + model=actual_model, + stopReason=stop_reason, + ) + # Simple text response + text = message.content or "" + return CreateMessageResult( + role="assistant", + content=TextContent(type="text", text=text), + model=actual_model, + stopReason=stop_reason, + ) + + +async def _check_model_access( # noqa: PLR0915 + model: str, user_api_key_auth: Any +) -> Optional["ErrorData"]: + """Enforce model-permission checks for MCP sampling requests. + + Runs the same authorization checks as ``/chat/completions``: + key-level, team-level, per-member, user-level, and project-level + model restrictions. The model name comes from the upstream MCP + server (untrusted input). + + Returns None if authorized, or an ErrorData describing the denial. + """ + if user_api_key_auth is None: + return None + + _api_key = getattr(user_api_key_auth, "api_key", None) + _token = getattr(user_api_key_auth, "token", None) + _user_role = getattr(user_api_key_auth, "user_role", None) + + _has_real_credential = bool(_api_key) or bool(_token) + _is_admin = ( + _user_role in ("proxy_admin", "proxy_admin_viewer") if _user_role else False + ) + + if not _has_real_credential and not _is_admin: + verbose_logger.warning( + "MCP sampling: denying model access for model=%s — " + "auth context has no real LiteLLM credential (possible " + "OAuth passthrough placeholder). api_key=%s, token=%s, role=%s", + model, + bool(_api_key), + bool(_token), + _user_role, + ) + return ErrorData( + code=-1, + message=( + "Model access denied: sampling requires a valid LiteLLM " + "API key or admin credential. OAuth-only sessions cannot " + "trigger proxy model calls without explicit authorization." + ), + ) + + try: + import litellm + from litellm.proxy.auth.auth_checks import ( + can_key_call_model, + can_team_access_model, + can_user_call_model, + can_project_access_model, + _check_team_member_model_access, + get_team_object, + get_user_object, + get_project_object, + ) + + try: + from litellm.proxy.proxy_server import llm_router as _llm_router + except ImportError: + _llm_router = None + + await can_key_call_model( + model=model, + llm_model_list=getattr(litellm, "model_list", None), + valid_token=user_api_key_auth, + llm_router=_llm_router, + ) + + _team_id = getattr(user_api_key_auth, "team_id", None) + _user_id = getattr(user_api_key_auth, "user_id", None) + _project_id = getattr(user_api_key_auth, "project_id", None) + + try: + from litellm.proxy.proxy_server import ( + prisma_client as _prisma_client, + user_api_key_cache as _user_api_key_cache, + proxy_logging_obj as _proxy_logging_obj, + ) + except ImportError: + _prisma_client = None + _user_api_key_cache = None # type: ignore[assignment] + _proxy_logging_obj = None # type: ignore[assignment] + + if _team_id and _prisma_client and _user_api_key_cache: + try: + team_obj = await get_team_object( + team_id=_team_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + team_obj = None + + if team_obj: + await can_team_access_model( + model=model, + team_object=team_obj, + llm_router=_llm_router, + team_model_aliases=getattr( + user_api_key_auth, "team_model_aliases", None + ), + ) + if _user_id and _proxy_logging_obj: + await _check_team_member_model_access( + model=model, + team_object=team_obj, + valid_token=user_api_key_auth, + llm_router=_llm_router, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + elif not _team_id and _user_id and _prisma_client and _user_api_key_cache: + try: + user_obj = await get_user_object( + user_id=_user_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + user_obj = None + + if user_obj: + await can_user_call_model( + model=model, + llm_router=_llm_router, + user_object=user_obj, + ) + + if _project_id and _prisma_client and _user_api_key_cache: + try: + project_obj = await get_project_object( + project_id=_project_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + project_obj = None + + if project_obj: + can_project_access_model( + model=model, + project_object=project_obj, + llm_router=_llm_router, + ) + + verbose_logger.debug( + "MCP sampling: model access check passed for model=%s", + model, + ) + return None + except Exception as access_err: + verbose_logger.warning( + "MCP sampling: model access denied for model=%s: %s", + model, + access_err, + ) + return ErrorData( + code=-1, + message=( + f"Model access denied: the API key is not authorized " + f"to use model '{model}'. {access_err}" + ), + ) + + +async def _run_budget_checks( + model: str, + user_api_key_auth: Any, + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, +) -> Optional["ErrorData"]: + """Enforce key/team/user/org/global budget checks for sampling requests. + + Runs the same ``common_checks`` path that ``/chat/completions`` uses, + so sampling cannot bypass budget limits. + + Returns None if all checks pass, or an ErrorData describing the denial. + """ + try: + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.proxy_server import ( + general_settings, + llm_router as _llm_router, + prisma_client as _prisma_client, + proxy_logging_obj as _proxy_logging_obj, + user_api_key_cache as _user_api_key_cache, + ) + from litellm.proxy.auth.auth_checks import ( + get_team_object, + get_user_object, + ) + import litellm + except ImportError as import_err: + verbose_logger.warning( + "MCP sampling: budget check imports unavailable: %s", import_err + ) + return None # Can't enforce budgets without the modules + + _team_id = getattr(user_api_key_auth, "team_id", None) + _user_id = getattr(user_api_key_auth, "user_id", None) + + team_obj = None + if _team_id and _prisma_client and _user_api_key_cache: + try: + team_obj = await get_team_object( + team_id=_team_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + pass + + user_obj = None + if _user_id and _prisma_client and _user_api_key_cache: + try: + user_obj = await get_user_object( + user_id=_user_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + pass + + dummy_request = _build_sampling_request( + raw_headers=raw_headers, + client_ip=client_ip, + ) + + # Enforce virtual-key route restrictions: a key limited to MCP routes + # must not be able to trigger a /chat/completions call via sampling. + # This mirrors the RouteChecks.should_call_route gate that runs in + # user_api_key_auth before common_checks for regular requests. + try: + from litellm.proxy.auth.route_checks import RouteChecks + + RouteChecks.should_call_route( + route="/chat/completions", + valid_token=user_api_key_auth, + request=dummy_request, + ) + except HTTPException as route_err: + verbose_logger.warning( + "MCP sampling: route check denied /chat/completions for key: %s", + route_err.detail, + ) + return ErrorData( + code=-1, + message=f"Sampling denied: virtual key is not allowed to call /chat/completions. {route_err.detail}", + ) + + global_proxy_spend = getattr(litellm, "_global_proxy_spend", None) + + # Build request body and merge x-litellm-tags from MCP headers BEFORE + # common_checks runs. _tag_max_budget_check inside common_checks only + # inspects request_body; without this pre-merge, header-supplied tags + # bypass per-tag budget enforcement (mirroring the regular auth path). + request_body: Dict[str, Any] = {"model": model} + try: + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=dummy_request, + request_data=request_body, + user_api_key_dict=user_api_key_auth, + ) + except Exception: + # Non-fatal: tag merge is defense-in-depth; don't block sampling + # if the merge utility is unavailable or fails. + pass + + try: + await common_checks( + request_body=request_body, + team_object=team_obj, + user_object=user_obj, + end_user_object=None, + global_proxy_spend=global_proxy_spend, + general_settings=general_settings or {}, + route="/chat/completions", + llm_router=_llm_router, + proxy_logging_obj=typing.cast("ProxyLogging", _proxy_logging_obj), + valid_token=user_api_key_auth, + request=dummy_request, + ) + except Exception as budget_err: + verbose_logger.warning( + "MCP sampling: budget check failed for model=%s: %s", + model, + budget_err, + ) + return ErrorData( + code=-1, + message=f"Sampling denied: {budget_err}", + ) + + verbose_logger.debug("MCP sampling: budget checks passed for model=%s", model) + return None + + +def _build_sampling_request( + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, +) -> Any: + """Build a synthetic FastAPI Request for sampling sub-calls. + + Converts the original MCP connection's HTTP headers into ASGI + scope format so that ``add_litellm_data_to_request`` can apply + header-dependent guardrails, tag-based routing, trace correlation, + and ``forward_llm_provider_auth_headers``. + + Key fields populated: + - **headers**: All original HTTP headers are forwarded (except + hop-by-hop: content-length, transfer-encoding). This ensures + ``traceparent``, ``authorization``, ``user-agent``, and + ``x-litellm-api-key`` are visible to pre-call utils. + - **client**: The ASGI ``(host, port)`` tuple so that + ``request.client.host`` returns the real client IP for + IP-based routing and guardrails. + - **server**: Derived from the running proxy's ``server_host`` + / ``server_port`` when available, avoiding the misleading + ``127.0.0.1:0`` placeholder. + - **x-forwarded-for**: Injected from ``client_ip`` if the + original headers don't already carry it, as a fallback for + IP attribution. + """ + from fastapi import Request + + # --- Build ASGI headers --- + _scope_headers: list = [(b"content-type", b"application/json")] + # Hop-by-hop headers that must NOT be forwarded into the + # synthetic request (they describe the original HTTP framing, + # not the logical request). + _HOP_BY_HOP = frozenset( + { + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "te", + "trailer", + } + ) + if raw_headers: + for hdr_name, hdr_value in raw_headers.items(): + _key = hdr_name.lower() + # Skip content-type (already set), x-forwarded-for (use resolved + # client_ip instead to prevent spoofing), and hop-by-hop headers + if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP: + continue + _scope_headers.append( + ( + _key.encode("latin-1", errors="replace"), + hdr_value.encode("utf-8"), + ) + ) + + # Inject x-forwarded-for from captured client_ip if the + # original headers don't already carry it + if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers): + _scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8"))) + + # --- Derive server (host, port) from the running proxy --- + _server_host = "127.0.0.1" + _server_port = 4000 # LiteLLM default + try: + import litellm.proxy.proxy_server as proxy_server + + _proxy_host = getattr(proxy_server, "server_host", None) + _proxy_port = getattr(proxy_server, "server_port", None) + + if _proxy_host: + _server_host = str(_proxy_host) + if _proxy_port: + _server_port = int(_proxy_port) + except (ImportError, AttributeError, TypeError, ValueError): + pass + + # --- Build ASGI client tuple for request.client.host --- + _client_tuple = None + if client_ip: + _client_tuple = (client_ip, 0) + + scope: Dict[str, Any] = { + "type": "http", + "method": "POST", + "path": "/mcp/sampling/createMessage", + "scheme": "http", + "server": (_server_host, _server_port), + "query_string": b"", + "root_path": "", + "headers": _scope_headers, + } + if _client_tuple is not None: + scope["client"] = _client_tuple + + return Request(scope=scope) + + +async def _build_completion_kwargs( + params: "CreateMessageRequestParams", + model: str, + user_api_key_auth: Any, + raw_headers: Optional[Dict[str, str]], + client_ip: Optional[str], +) -> Dict[str, Any]: + openai_messages = _convert_mcp_messages_to_openai( + messages=params.messages, + system_prompt=params.systemPrompt, + ) + completion_kwargs: Dict[str, Any] = { + "model": model, + "messages": openai_messages, + "max_tokens": params.maxTokens, + } + if params.temperature is not None: + completion_kwargs["temperature"] = params.temperature + if params.stopSequences: + completion_kwargs["stop"] = params.stopSequences + openai_tools = _convert_mcp_tools_to_openai(params.tools) + if openai_tools: + completion_kwargs["tools"] = openai_tools + openai_tool_choice = _convert_mcp_tool_choice_to_openai(params.toolChoice) + if openai_tool_choice is not None: + completion_kwargs["tool_choice"] = openai_tool_choice + completion_kwargs["metadata"] = {} + if params.metadata: + completion_kwargs["metadata"]["mcp_metadata"] = params.metadata + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import proxy_config + + completion_kwargs["user"] = getattr(user_api_key_auth, "user_id", None) + _dummy_request = _build_sampling_request( + raw_headers=raw_headers, client_ip=client_ip + ) + completion_kwargs = await add_litellm_data_to_request( + data=completion_kwargs, + request=_dummy_request, + user_api_key_dict=user_api_key_auth, + proxy_config=proxy_config, + ) + return completion_kwargs + + +async def _run_guardrails_and_call_llm( + completion_kwargs: Dict[str, Any], + user_api_key_auth: Any, +) -> Any: + try: + from litellm.proxy.proxy_server import proxy_logging_obj as _plo + + if _plo is not None: + completion_kwargs = await typing.cast("ProxyLogging", _plo).pre_call_hook( + user_api_key_dict=user_api_key_auth, + data=completion_kwargs, + call_type="acompletion", + ) + except ImportError: + pass + except Exception as guardrail_err: + verbose_logger.warning( + "MCP sampling: pre-call guardrail rejected request: %s", + guardrail_err, + ) + raise + + import litellm + + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is not None: + return await llm_router.acompletion(**completion_kwargs) + return await litellm.acompletion(**completion_kwargs) + except ImportError: + return await litellm.acompletion(**completion_kwargs) + + +async def handle_sampling_create_message( + context: Any, + params: "CreateMessageRequestParams", + default_model: Optional[str] = None, + user_api_key_auth: Optional[Any] = None, + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, +) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: + """ + Handle an MCP sampling/createMessage request by routing through LiteLLM. + This is the main entry point called by the MCP client session when an + upstream MCP server requests LLM inference. + Args: + context: MCP RequestContext (contains session info). + params: The CreateMessageRequestParams from the MCP server. + default_model: Default model to use if no preferences match. + user_api_key_auth: Auth context for the requesting user. + raw_headers: Original HTTP headers from the MCP connection. + Forwarded into the internal acompletion call so that + header-dependent guardrails, IP-routing, trace-id + correlation, and forward_llm_provider_auth_headers + work correctly for sampling sub-calls. + client_ip: Original client IP address for IP-based guardrails. + Returns: + CreateMessageResult with the LLM's response, or ErrorData on failure. + """ + if not MCP_SAMPLING_AVAILABLE: + return ErrorData( + code=-1, + message="MCP sampling is not available (mcp package not installed)", + ) + + if user_api_key_auth is None: + return ErrorData( + code=-1, + message=( + "Sampling requires an authenticated user context. " + "Internal or unauthenticated sessions cannot trigger " + "upstream-initiated model calls." + ), + ) + + try: + model = _resolve_model_from_preferences( + model_preferences=params.modelPreferences, + default_model=default_model, + ) + verbose_logger.info( + "MCP sampling: resolved model=%s from preferences=%s", + model, + params.modelPreferences, + ) + + access_denial = await _check_model_access(model, user_api_key_auth) + if access_denial is not None: + return access_denial + + budget_denial = await _run_budget_checks( + model=model, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if budget_denial is not None: + return budget_denial + + completion_kwargs = await _build_completion_kwargs( + params=params, + model=model, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + + openai_messages = completion_kwargs["messages"] + openai_tools = completion_kwargs.get("tools") + verbose_logger.debug( + "MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s", + model, + len(openai_messages), + bool(openai_tools), + ) + + response = await _run_guardrails_and_call_llm( + completion_kwargs=completion_kwargs, + user_api_key_auth=user_api_key_auth, + ) + + result = _convert_openai_response_to_mcp_result( + response=response, model_name=model + ) + verbose_logger.info( + "MCP sampling: completed successfully, model=%s, stopReason=%s", + getattr(result, "model", "unknown"), + getattr(result, "stopReason", "unknown"), + ) + return result + except Exception as e: + from litellm.exceptions import ( + AuthenticationError, + BudgetExceededError, + ContextWindowExceededError, + PermissionDeniedError, + RateLimitError, + ServiceUnavailableError, + ) + + from litellm.proxy._types import ProxyException + + if isinstance( + e, + ( + HTTPException, + BudgetExceededError, + RateLimitError, + AuthenticationError, + PermissionDeniedError, + ContextWindowExceededError, + ServiceUnavailableError, + ProxyException, + ), + ): + raise + + verbose_logger.exception("MCP sampling handler failed: %s", e) + return ErrorData( + code=-1, + message=f"Sampling failed: {str(e)}", + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6e33a105ec8..df6cb22fda1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,6 +6,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib +import contextvars import hashlib import json import time @@ -125,6 +126,18 @@ try: GetPromptResult, ResourceTemplate, TextResourceContents, + Tool, + ) + from mcp.server.session import ServerSession as _McpServerSession + import weakref + + # Robust auth lookup keyed by session_object. + _session_obj_auth_storage: ( + "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" + ) = weakref.WeakKeyDictionary() + + active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = ( + contextvars.ContextVar("active_mcp_session", default=None) ) except ImportError as e: verbose_logger.debug(f"MCP module not found: {e}") @@ -160,6 +173,60 @@ def _mcp_session_id_from_headers( return None +def _jsonrpc_text_has_top_level_method(text: str) -> bool: + """Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at + the root object's top level. + + Used to tell a request/notification (carries ``method``) apart from a + response (carries ``result``/``error`` and no top-level ``method``). A + response payload can itself nest a ``method`` field, so only keys at the + root object's depth are inspected rather than searching the whole string. + Returns ``True`` only when a top-level ``method`` key is positively found; + truncation that hides it yields ``False``. + """ + depth = 0 + in_string = False + escaped = False + in_object: List[bool] = [] + reading_key = False + expect_key = False + key_chars: List[str] = [] + for ch in text: + if in_string: + if escaped: + escaped = False + elif ch == "\\": + escaped = True + elif ch == '"': + in_string = False + if reading_key and depth == 1 and "".join(key_chars) == "method": + return True + elif reading_key: + key_chars.append(ch) + continue + if ch == '"': + in_string = True + reading_key = expect_key and depth >= 1 and in_object[-1] + key_chars = [] + expect_key = False + elif ch == "{" or ch == "[": + depth += 1 + in_object.append(ch == "{") + expect_key = ch == "{" + elif ch == "}" or ch == "]": + if in_object: + in_object.pop() + depth -= 1 + if depth <= 0: + break + expect_key = False + elif ch == ",": + expect_key = bool(in_object) and in_object[-1] + elif ch == ":": + expect_key = False + return False + + if MCP_AVAILABLE: from mcp.server import Server from mcp.server.lowlevel.server import NotificationOptions @@ -483,10 +550,18 @@ if MCP_AVAILABLE: ######################################################## @server.list_tools() - async def list_tools() -> List[MCPTool]: + async def handle_list_tools() -> List[Tool]: """ - List all available tools + List all available tools. + Also captures the active session for propagation to callbacks. """ + from mcp.server.lowlevel.server import request_ctx + + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) + try: # Get user authentication from context variable ( @@ -497,7 +572,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}" ) @@ -528,152 +603,178 @@ if MCP_AVAILABLE: # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) @server.call_tool() - async def mcp_server_tool_call( - name: str, arguments: Optional[Dict[str, Any]] + async def mcp_server_tool_call( # noqa: PLR0915 + name: str, arguments: Dict[str, Any] | None ) -> CallToolResult: """ Call a specific tool with the provided arguments - Args: name (str): Name of the tool to call arguments (Dict[str, Any] | None): Arguments to pass to the tool - Returns: List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results - Raises: HTTPException: If tool not found or arguments missing """ from fastapi import Request - from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import proxy_config + from mcp.types import CallToolResult + from mcp.server.lowlevel.server import request_ctx - # Validate arguments - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = get_auth_context() + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) - verbose_logger.debug( - f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" - ) - host_progress_callback = None try: - host_ctx = server.request_context - if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta: - host_token = getattr(host_ctx.meta, "progressToken", None) - if host_token and hasattr(host_ctx, "session") and host_ctx.session: - host_session = host_ctx.session - - async def forward_progress(progress: float, total: Optional[float]): - """Forward progress notifications from external MCP to Host""" - try: - await host_session.send_progress_notification( - progress_token=host_token, - progress=progress, - total=total, - ) - verbose_logger.debug( - f"Forwarded progress {progress}/{total} to Host" - ) - except Exception as e: - verbose_logger.error( - f"Failed to forward progress to Host: {e}" - ) - - host_progress_callback = forward_progress - verbose_logger.debug( - f"Host progressToken captured: {host_token[:8]}..." - ) - except Exception as e: - verbose_logger.warning(f"Could not capture host progress context: {e}") - try: - # Create a body date for logging - body_data = {"name": name, "arguments": arguments} - # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) - chain_id = get_chain_id_from_headers(raw_headers) - if chain_id: - body_data["litellm_trace_id"] = chain_id - body_data["litellm_session_id"] = chain_id - - request = Request( - scope={ - "type": "http", - "method": "POST", - "path": "/mcp/tools/call", - "headers": [(b"content-type", b"application/json")], - } + # Validate arguments + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = await get_or_extract_auth_context() + verbose_logger.debug( + f"MCP mcp_server_tool_call - user_api_key_auth={user_api_key_auth}, user_role={getattr(user_api_key_auth, 'user_role', 'N/A')}" ) - if user_api_key_auth is not None: - data = await add_litellm_data_to_request( - data=body_data, - request=request, - user_api_key_dict=user_api_key_auth, - proxy_config=proxy_config, + + verbose_logger.debug( + f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" + ) + host_progress_callback = None + try: + host_ctx = server.request_context + if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta: + host_token = getattr(host_ctx.meta, "progressToken", None) + if host_token and hasattr(host_ctx, "session") and host_ctx.session: + host_session = host_ctx.session + + async def forward_progress( + progress: float, total: Optional[float] + ): + """Forward progress notifications from external MCP to Host""" + try: + await host_session.send_progress_notification( + progress_token=host_token, + progress=progress, + total=total, + ) + verbose_logger.debug( + f"Forwarded progress {progress}/{total} to Host" + ) + except Exception as e: + verbose_logger.error( + f"Failed to forward progress to Host: {e}" + ) + + host_progress_callback = forward_progress + verbose_logger.debug( + f"Host progressToken captured: {host_token[:8]}..." + ) + except Exception as e: + verbose_logger.warning(f"Could not capture host progress context: {e}") + try: + # Create a body date for logging + body_data = {"name": name, "arguments": arguments} + # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) + chain_id = get_chain_id_from_headers(raw_headers) + if chain_id: + body_data["litellm_trace_id"] = chain_id + body_data["litellm_session_id"] = chain_id + + request = Request( + scope={ + "type": "http", + "method": "POST", + "path": "/mcp/tools/call", + "headers": [(b"content-type", b"application/json")], + } ) - else: - data = body_data - - response = await call_mcp_tool( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - host_progress_callback=host_progress_callback, - **data, # for logging - ) - except BlockedPiiEntityError as e: - verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}") - return CallToolResult( - content=[ - TextContent( - text=f"Error: Blocked PII entity detected - {str(e)}", - type="text", + if user_api_key_auth is not None: + data = await add_litellm_data_to_request( + data=body_data, + request=request, + user_api_key_dict=user_api_key_auth, + proxy_config=proxy_config, ) - ], - isError=True, - ) - except GuardrailRaisedException as e: - verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}") - return CallToolResult( - content=[ - TextContent( - text=f"Error: Guardrail violation - {str(e)}", type="text" - ) - ], - isError=True, - ) - except HTTPException as e: - verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") - return CallToolResult( - content=[TextContent(text=f"Error: {str(e.detail)}", type="text")], - isError=True, - ) - except Exception as e: - verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}") - return CallToolResult( - content=[TextContent(text=f"Error: {str(e)}", type="text")], - isError=True, - ) + else: + data = body_data - return response + response = await call_mcp_tool( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + host_progress_callback=host_progress_callback, + **data, # for logging + ) + except BlockedPiiEntityError as e: + verbose_logger.error( + f"BlockedPiiEntityError in MCP tool call: {str(e)}" + ) + return CallToolResult( + content=[ + TextContent( + text=f"Error: Blocked PII entity detected - {str(e)}", + type="text", + ) + ], + isError=True, + ) + except GuardrailRaisedException as e: + verbose_logger.error( + f"GuardrailRaisedException in MCP tool call: {str(e)}" + ) + return CallToolResult( + content=[ + TextContent( + text=f"Error: Guardrail violation - {str(e)}", type="text" + ) + ], + isError=True, + ) + except HTTPException as e: + verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") + return CallToolResult( + content=[TextContent(text=f"Error: {str(e.detail)}", type="text")], + isError=True, + ) + except Exception as e: + verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}") + return CallToolResult( + content=[TextContent(text=f"Error: {str(e)}", type="text")], + isError=True, + ) + + return response + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) @server.list_prompts() async def list_prompts() -> List[Prompt]: """ List all available prompts """ + from mcp.server.lowlevel.server import request_ctx + + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) + try: # Get user authentication from context variable ( @@ -684,7 +785,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}" ) @@ -713,6 +814,9 @@ if MCP_AVAILABLE: # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) @server.get_prompt() async def get_prompt( @@ -730,33 +834,13 @@ if MCP_AVAILABLE: """ # Validate arguments - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = get_auth_context() + from mcp.server.lowlevel.server import request_ctx - verbose_logger.debug( - f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" - ) - return await mcp_get_prompt( - name=name, - arguments=arguments, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) - @server.list_resources() - async def list_resources() -> List[Resource]: - """List all available resources.""" try: ( user_api_key_auth, @@ -766,7 +850,45 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() + + verbose_logger.debug( + f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" + ) + return await mcp_get_prompt( + name=name, + arguments=arguments, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) + + @server.list_resources() + async def list_resources() -> List[Resource]: + """List all available resources.""" + from mcp.server.lowlevel.server import request_ctx + + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) + + try: + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}" ) @@ -792,10 +914,20 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.exception(f"Error in list_resources endpoint: {str(e)}") return [] + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) @server.list_resource_templates() async def list_resource_templates() -> List[ResourceTemplate]: """List all available resource templates.""" + from mcp.server.lowlevel.server import request_ctx + + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) + try: ( user_api_key_auth, @@ -805,7 +937,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}" ) @@ -825,8 +957,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) verbose_logger.info( - "MCP list_resource_templates - Successfully returned " - f"{len(resource_templates)} resource templates" + f"MCP list_resource_templates - Successfully returned {len(resource_templates)} resource templates" ) return resource_templates except Exception as e: @@ -834,30 +965,44 @@ if MCP_AVAILABLE: f"Error in list_resource_templates endpoint: {str(e)}" ) return [] + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) @server.read_resource() async def read_resource(url: AnyUrl) -> list[ReadResourceContents]: - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = get_auth_context() + from mcp.server.lowlevel.server import request_ctx - read_resource_result = await mcp_read_resource( - url=url, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) + req_ctx = request_ctx.get(None) + _session_reset_token = None + if req_ctx: + _session_reset_token = active_mcp_session_var.set(req_ctx.session) - return _normalize_resource_contents(read_resource_result.contents) + try: + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = await get_or_extract_auth_context() + + read_resource_result = await mcp_read_resource( + url=url, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + + return _normalize_resource_contents(read_resource_result.contents) + finally: + if _session_reset_token is not None: + active_mcp_session_var.reset(_session_reset_token) ######################################################## ############ End of MCP Server Routes ################## @@ -1180,8 +1325,7 @@ if MCP_AVAILABLE: cached_token = await mcp_per_user_token_cache.get(user_id, server_id) if cached_token is not None: verbose_logger.debug( - "_get_user_oauth_extra_headers_from_db: Redis hit for " - "user=%s server=%s", + "_get_user_oauth_extra_headers_from_db: Redis hit for user=%s server=%s", user_id, server_id, ) @@ -1207,8 +1351,7 @@ if MCP_AVAILABLE: if is_oauth_credential_expired(cred): verbose_logger.debug( - "_get_user_oauth_extra_headers_from_db: token expired for " - "user=%s server=%s — attempting refresh", + "_get_user_oauth_extra_headers_from_db: token expired for user=%s server=%s — attempting refresh", user_id, server_id, ) @@ -1230,8 +1373,7 @@ if MCP_AVAILABLE: ) except Exception as refresh_exc: verbose_logger.warning( - "_get_user_oauth_extra_headers_from_db: refresh failed " - "for user=%s server=%s: %s", + "_get_user_oauth_extra_headers_from_db: refresh failed for user=%s server=%s: %s", user_id, server_id, refresh_exc, @@ -1275,8 +1417,7 @@ if MCP_AVAILABLE: return {"Authorization": f"Bearer {access_token}"} except Exception as e: verbose_logger.warning( - "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " - "user=%s server=%s: %s", + "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for user=%s server=%s: %s", user_id, server_id, e, @@ -2485,7 +2626,7 @@ if MCP_AVAILABLE: arguments=arguments or {}, server_name=server_name or mcp_server.name, user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=proxy_logging_obj, # type: ignore[arg-type] server=mcp_server, raw_headers=raw_headers, ) @@ -2744,8 +2885,7 @@ if MCP_AVAILABLE: raise HTTPException( status_code=400, detail=( - "Multiple MCP servers configured; read_resource currently " - "supports exactly one allowed server." + "Multiple MCP servers configured; read_resource currently supports exactly one allowed server." ), ) @@ -3124,8 +3264,7 @@ if MCP_AVAILABLE: return False except Exception: verbose_logger.debug( - "Unable to inspect active MCP sessions for '%s'. " - "Deferring to session manager.", + "Unable to inspect active MCP sessions for '%s'. Deferring to session manager.", _session_id, ) return False @@ -3136,8 +3275,7 @@ if MCP_AVAILABLE: if method == "DELETE": _remove_stateful_session_tracking(_session_id) verbose_logger.info( - "DELETE request for non-existent MCP session '%s'. " - "Returning success (idempotent DELETE).", + "DELETE request for non-existent MCP session '%s'. Returning success (idempotent DELETE).", _session_id, ) success_response = JSONResponse( @@ -3615,6 +3753,7 @@ if MCP_AVAILABLE: return session_id = _get_session_id_from_scope(scope) + body = b"" if scope.get("method") == "POST": consumed_messages, body = await _read_request_body_for_routing(receive) is_initialize = _is_initialize_request(body) @@ -3639,8 +3778,7 @@ if MCP_AVAILABLE: ) if not await _enforce_stateful_session_cap_for_owner(request_owner): verbose_logger.warning( - "Rejecting MCP initialize: caller already holds the maximum " - "number of active stateful sessions." + "Rejecting MCP initialize: caller already holds the maximum number of active stateful sessions." ) too_many_response = JSONResponse( status_code=429, @@ -3672,9 +3810,56 @@ if MCP_AVAILABLE: # POST/DELETE are the methods that actually mutate the shared # auth context, so serializing those is sufficient for the # clobbering race between concurrent JSON-RPC calls. - session_lock: Optional[asyncio.Lock] = None + # + # Also skip the lock for JSON-RPC *responses* (POSTs that carry + # a ``result`` or ``error`` but no ``method``). These are replies + # to server-initiated requests such as ``elicitation/create`` or + # ``sampling/createMessage``. The in-flight tool-call POST that + # triggered the server request already holds the session lock, so + # trying to acquire it again for the response POST would deadlock. + is_jsonrpc_response = False request_method = (scope.get("method") or "").upper() - if use_stateful and session_id and request_method in ("POST", "DELETE"): + if body and request_method == "POST": + try: + _peeked = json.loads(body) + if ( + isinstance(_peeked, dict) + and _peeked.get("jsonrpc") == "2.0" + and "id" in _peeked + and "method" not in _peeked + and ("result" in _peeked or "error" in _peeked) + ): + is_jsonrpc_response = True + verbose_logger.debug( + "MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock", + _peeked.get("id"), + ) + except (json.JSONDecodeError, TypeError): + # Peek cap truncated the body, so it can't be fully parsed. + # Scan the top-level keys (depth-aware) instead of a flat + # substring search: a response's result payload may nest a + # "method" field, and misreading that would acquire the lock + # and deadlock the in-flight tool call awaiting this + # response. A false skip is harmless; a false acquire is not. + _body_str = body.decode("utf-8", errors="replace") + if ( + '"jsonrpc"' in _body_str + and ('"result"' in _body_str or '"error"' in _body_str) + and not _jsonrpc_text_has_top_level_method(_body_str) + ): + is_jsonrpc_response = True + verbose_logger.debug( + "MCP: detected truncated JSON-RPC response POST via " + "top-level key scan, skipping session lock to avoid deadlock" + ) + + session_lock: Optional[asyncio.Lock] = None + if ( + use_stateful + and session_id + and request_method in ("POST", "DELETE") + and not is_jsonrpc_response + ): session_lock = _stateful_session_locks.setdefault( session_id, asyncio.Lock() ) @@ -4099,6 +4284,119 @@ if MCP_AVAILABLE: ) return None, None, None, None, None, None, None + def _get_current_session(): + try: + from mcp.server.lowlevel.server import request_ctx + + return request_ctx.get().session + except (LookupError, ImportError): + return None + + def _cache_auth_context_lazily(): + session = _get_current_session() + if session is None: + return + try: + if session in _session_obj_auth_storage: + return + except TypeError: + verbose_logger.debug( + "_cache_auth_context_lazily: session object is unhashable (type=%s), cannot cache auth context", + type(session).__name__, + ) + return + + auth = auth_context_var.get() + if auth and isinstance(auth, MCPAuthenticatedUser): + try: + _session_obj_auth_storage[session] = auth + except TypeError: + verbose_logger.debug( + "_cache_auth_context_lazily: could not store auth via " + "session identity — session object is unhashable" + ) + + def _recover_auth_from_session() -> Optional[MCPAuthenticatedUser]: + session = _get_current_session() + if session is None: + return None + + stored: Optional[MCPAuthenticatedUser] = None + try: + stored = _session_obj_auth_storage.get(session) + except TypeError: + verbose_logger.debug( + "_recover_auth_from_session: session object is unhashable " + "(type=%s), skipping _session_obj_auth_storage lookup", + type(session).__name__, + ) + + return stored + + async def get_or_extract_auth_context() -> Tuple[ + Optional[UserAPIKeyAuth], + Optional[str], + Optional[List[str]], + Optional[Dict[str, Dict[str, str]]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], + Optional[str], + ]: + """ + Get auth context from ContextVar first, then fall back to session + storage (which survives cross-task boundaries in the MCP SDK). + """ + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = get_auth_context() + + if user_api_key_auth is not None: + _cache_auth_context_lazily() + else: + stored = _recover_auth_from_session() + + if stored: + user_api_key_auth = stored.user_api_key_auth + mcp_auth_header = stored.mcp_auth_header + mcp_servers = stored.mcp_servers + mcp_server_auth_headers = stored.mcp_server_auth_headers + oauth2_headers = stored.oauth2_headers + raw_headers = stored.raw_headers + _client_ip = stored.client_ip + return ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) + + def get_active_mcp_session() -> Optional[_McpServerSession]: + """Return the active MCP session captured during handler execution.""" + session = active_mcp_session_var.get() + if session is not None: + return session + return _get_current_session() + + def get_active_auth_context() -> Optional[MCPAuthenticatedUser]: + """Return auth context from ContextVar or session storage.""" + auth = auth_context_var.get() + if auth and isinstance(auth, MCPAuthenticatedUser): + return auth + + stored = _recover_auth_from_session() + if stored is not None: + return stored + return None + ######################################################## ############ End of Auth Context Functions ############# ######################################################## diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 20c951d350d..f0d8ddf97d6 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -60,11 +60,20 @@ async def cache_ping(): """ litellm_cache_params: Dict[str, Any] = {} cleaned_cache_params: Dict[str, Any] = {} + if litellm.cache is None: + raise ProxyException( + message=safe_dumps( + { + "message": "Cache not initialized. litellm.cache is None", + "litellm_cache_params": "{}", + "health_check_cache_params": "{}", + } + ), + type=ProxyErrorTypes.cache_ping_error, + param="cache_ping", + code=503, + ) try: - if litellm.cache is None: - raise HTTPException( - status_code=503, detail="Cache not initialized. litellm.cache is None" - ) litellm_cache_params = masker.mask_dict(vars(litellm.cache)) # remove field that might reference itself litellm_cache_params.pop("cache", None) @@ -97,14 +106,14 @@ async def cache_ping(): cache_type=str(litellm.cache.type), litellm_cache_params=safe_dumps(litellm_cache_params), ) - except Exception as e: - import traceback - + except HTTPException: + raise + except Exception: + verbose_proxy_logger.exception("Cache health check failed") error_message = { - "message": f"Service Unhealthy ({str(e)})", + "message": "Service Unhealthy", "litellm_cache_params": safe_dumps(litellm_cache_params), "health_check_cache_params": safe_dumps(cleaned_cache_params), - "traceback": traceback.format_exc(), } raise ProxyException( message=safe_dumps(error_message), diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 27fa685eaac..b0932015ab3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -16,7 +16,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ToolPermissionRule, @@ -60,53 +60,7 @@ class ToolPermissionGuardrail(CustomGuardrail): super().__init__(**kwargs) - self.rules: List[ToolPermissionRule] = [] - self._compiled_rule_patterns: Dict[str, Dict[str, re.Pattern]] = {} - self._compiled_rule_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {} - if rules: - for rule_item in rules: - if isinstance(rule_item, ToolPermissionRule): - rule = rule_item - else: - rule = ToolPermissionRule(**rule_item) - self.rules.append(rule) - - compiled_target_patterns: Dict[str, Optional[re.Pattern]] = { - "tool_name": None, - "tool_type": None, - } - if rule.tool_name is not None: - try: - compiled_target_patterns["tool_name"] = re.compile( - rule.tool_name - ) - except re.error as exc: - raise ValueError( - f"Invalid regex for tool_name in rule '{rule.id}': {exc}" - ) from exc - if rule.tool_type is not None: - try: - compiled_target_patterns["tool_type"] = re.compile( - rule.tool_type - ) - except re.error as exc: - raise ValueError( - f"Invalid regex for tool_type in rule '{rule.id}': {exc}" - ) from exc - self._compiled_rule_targets[rule.id] = compiled_target_patterns - - if rule.allowed_param_patterns: - compiled_patterns: Dict[str, re.Pattern] = {} - for path, pattern in rule.allowed_param_patterns.items(): - try: - compiled_patterns[path] = re.compile(pattern) - except re.error as exc: - raise ValueError( - f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}" - ) from exc - - if compiled_patterns: - self._compiled_rule_patterns[rule.id] = compiled_patterns + self._load_rules(rules) # Normalize to lowercase for case-insensitive handling self.default_action = ( @@ -126,6 +80,115 @@ class ToolPermissionGuardrail(CustomGuardrail): self.default_action, ) + def _load_rules(self, rules: Optional[List[Any]]) -> None: + """Parse ``rules`` and (re)build the compiled target/pattern lookups. + + ``self.rules`` plus ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` + are the state every matching path reads. Centralizing the build here lets + both ``__init__`` and ``update_in_memory_litellm_params`` recompile from a + single source of truth, so an in-place update (PUT /guardrails, immediate + sync) reflects rule changes instead of keeping the construction-time maps. + """ + parsed_rules: List[ToolPermissionRule] = [] + compiled_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {} + compiled_patterns: Dict[str, Dict[str, re.Pattern]] = {} + + for rule_item in rules or []: + rule = ( + rule_item + if isinstance(rule_item, ToolPermissionRule) + else ToolPermissionRule(**rule_item) + ) + + target_patterns: Dict[str, Optional[re.Pattern]] = { + "tool_name": None, + "tool_type": None, + } + if rule.tool_name is not None: + try: + target_patterns["tool_name"] = re.compile(rule.tool_name) + except re.error as exc: + raise ValueError( + f"Invalid regex for tool_name in rule '{rule.id}': {exc}" + ) from exc + if rule.tool_type is not None: + try: + target_patterns["tool_type"] = re.compile(rule.tool_type) + except re.error as exc: + raise ValueError( + f"Invalid regex for tool_type in rule '{rule.id}': {exc}" + ) from exc + + rule_patterns: Dict[str, re.Pattern] = {} + for path, pattern in (rule.allowed_param_patterns or {}).items(): + try: + rule_patterns[path] = re.compile(pattern) + except re.error as exc: + raise ValueError( + f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}" + ) from exc + + parsed_rules.append(rule) + compiled_targets[rule.id] = target_patterns + if rule_patterns: + compiled_patterns[rule.id] = rule_patterns + + # Swap in the fully-built maps only after every rule compiles, so an + # invalid regex raises without leaving a partially-built ruleset (a + # missing compiled target is read as a match-all wildcard). + self.rules = parsed_rules + self._compiled_rule_targets = compiled_targets + self._compiled_rule_patterns = compiled_patterns + + def update_in_memory_litellm_params( + self, litellm_params: Union[LitellmParams, dict] + ) -> None: + """Apply updated params in place, rebuilding the compiled rule state. + + The base implementation only ``setattr``s raw fields, which would leave + ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` (built in + ``__init__``) stale, so a guardrail updated without reinitialization would + keep enforcing the old ruleset. Recompile here so PUT /guardrails and the + immediate in-memory sync take effect, mirroring the PresidioGuardrail + override of this method. + """ + # ``litellm_params`` may arrive as the raw DB dict (the proxy ``cast()``s + # it to ``LitellmParams`` without converting), so handle both shapes. The + # base ``setattr`` loop is model-only, so apply the dict case here. + previous_rules = self.rules + if isinstance(litellm_params, dict): + params = litellm_params + for key, value in params.items(): + setattr(self, key, value) + else: + super().update_in_memory_litellm_params(litellm_params) + params = vars(litellm_params) + + # The generic update above sets ``self.rules`` from the incoming value + # (None on a partial update that omits rules), but never rebuilds the + # compiled maps. Rebuild them when rules are provided; otherwise restore + # the previous ruleset so a partial update doesn't silently wipe it. An + # explicit empty list still clears the rules. + rules = params.get("rules") + if rules is not None: + try: + self._load_rules(rules) + except Exception: + # The generic update above may have overwritten self.rules with + # the raw payload; restore the prior consistent ruleset so a + # rejected update can't leave the live guardrail enforcing a + # broken policy. + self.rules = previous_rules + raise + else: + self.rules = previous_rules + default_action = params.get("default_action") + if isinstance(default_action, str): + self.default_action = default_action.lower() + on_disallowed_action = params.get("on_disallowed_action") + if isinstance(on_disallowed_action, str): + self.on_disallowed_action = on_disallowed_action.lower() + @staticmethod def get_config_model(): from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a092df6cdf0..7aed9ad894a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1910,7 +1910,7 @@ prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = N store_model_in_db: bool = False open_telemetry_logger: Optional[OpenTelemetry] = None ### INITIALIZE GLOBAL LOGGING OBJECT ### -proxy_logging_obj = ProxyLogging( +proxy_logging_obj: ProxyLogging = ProxyLogging( user_api_key_cache=user_api_key_cache, premium_user=premium_user ) ### REDIS QUEUE ### @@ -15844,10 +15844,10 @@ async def toolset_mcp_route(toolset_name: str, request: Request): except HTTPException as e: raise e except Exception as e: - verbose_proxy_logger.error( - f"Error handling toolset MCP route for {toolset_name}: {str(e)}" + verbose_proxy_logger.exception( + "Error handling toolset MCP route for %s: %s", toolset_name, str(e) ) - raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") + raise HTTPException(status_code=500, detail="Internal server error") async def _mcp_forward_as_path(path_segment: str, request: Request): @@ -16028,7 +16028,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): except HTTPException as e: raise e except Exception as e: - verbose_proxy_logger.error( - f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}" + verbose_proxy_logger.exception( + "Error handling dynamic MCP route for %s: %s", mcp_server_name, str(e) ) - raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") + raise HTTPException(status_code=500, detail="Internal server error") diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index 38e6d533449..e24eb4aebb5 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -1,7 +1,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional -from typing_extensions import TypedDict +from typing_extensions import Required, TypedDict from .vertex_ai import ( GenerationConfig, @@ -233,10 +233,11 @@ class GeminiImageGenerationResponse(TypedDict): # Video Generation Types -class GeminiVideoGenerationInstance(TypedDict): +class GeminiVideoGenerationInstance(TypedDict, total=False): """Instance data for Gemini video generation request""" - prompt: str + prompt: Required[str] + image: Dict[str, Any] class GeminiVideoGenerationParameters(BaseModel): @@ -264,11 +265,6 @@ class GeminiVideoGenerationParameters(BaseModel): negativePrompt: Optional[str] = None """Text describing what not to include in the video.""" - image: Optional[Any] = None - """ - An initial image to animate (Image object). - """ - lastFrame: Optional[Any] = None """ The final image for interpolation video to transition. diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 6aa62c35106..2108fe8990d 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -115,6 +115,8 @@ class MCPServer(BaseModel): # different ``server_id`` values are bumped deterministically. Left # ``None`` in default-prefix mode. short_prefix: Optional[str] = None + allow_sampling: bool = False + allow_elicitation: bool = False model_config = ConfigDict(arbitrary_types_allowed=True) @property diff --git a/litellm/utils.py b/litellm/utils.py index 7cac830b2c2..d010391229b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4871,7 +4871,7 @@ def add_provider_specific_params_to_optional_params( ) is False ): - extra_body = passed_params.pop("extra_body", None) or {} + extra_body = dict(passed_params.pop("extra_body", None) or {}) for k in passed_params.keys(): if k not in openai_params and passed_params[k] is not None: extra_body[k] = passed_params[k] @@ -8909,6 +8909,16 @@ class ProviderConfigManager: return litellm.OpenRouterResponsesAPIConfig() elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMResponsesAPIConfig() + elif litellm.LlmProviders.BEDROCK_MANTLE == provider: + # Only OpenAI gpt frontier models (gpt-5.x, and future gpt-6 etc.) are + # served on the /openai/v1/responses path. gpt-oss and every non-OpenAI + # model on Mantle (nvidia, mistral, google, zai, ...) are chat-completions + # only and 400 on that path, so they fall through to None to keep the + # chat-completions emulation (see litellm/responses/main.py "config is None"). + model_lower = model.lower() if model else "" + if "openai.gpt-" in model_lower and "gpt-oss" not in model_lower: + return litellm.BedrockMantleResponsesAPIConfig() + return None return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ed6de4fa6b7..4c227656e5f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30424,21 +30424,32 @@ "supports_reasoning": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "snowflake/claude-3-5-sonnet": { + "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", - "max_input_tokens": 18000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", - "supports_computer_use": true + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.0000003, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true }, - "snowflake/deepseek-r1": { + "snowflake/deepseek-r1": { "litellm_provider": "snowflake", - "max_input_tokens": 32768, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", - "supports_reasoning": true + "input_cost_per_token": 0.00000135, + "output_cost_per_token": 0.0000054, + "supports_reasoning": true, + "supports_system_messages": true }, "snowflake/gemma-7b": { "litellm_provider": "snowflake", @@ -30492,23 +30503,34 @@ "snowflake/llama3.1-405b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat" + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "input_cost_per_token": 0.0000012, + "output_cost_per_token": 0.0000012, + "supports_function_calling": true, + "supports_system_messages": true }, "snowflake/llama3.1-70b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat" + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "input_cost_per_token": 0.00000072, + "output_cost_per_token": 0.00000072, + "supports_function_calling": true, + "supports_system_messages": true }, "snowflake/llama3.1-8b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat" + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "input_cost_per_token": 0.00000024, + "output_cost_per_token": 0.00000024, + "supports_system_messages": true }, "snowflake/llama3.2-1b": { "litellm_provider": "snowflake", @@ -30524,13 +30546,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/llama3.3-70b": { - "litellm_provider": "snowflake", + "snowflake/llama3.3-70b": { + "max_tokens": 16384, "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat" - }, + "max_output_tokens": 16384, + "input_cost_per_token": 0.00000072, + "output_cost_per_token": 0.00000072, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, "snowflake/mistral-7b": { "litellm_provider": "snowflake", "max_input_tokens": 32000, @@ -30545,12 +30571,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/mistral-large2": { + "snowflake/mistral-large2": { "litellm_provider": "snowflake", "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat" + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "input_cost_per_token": 0.000002, + "output_cost_per_token": 0.000006, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true }, "snowflake/mixtral-8x7b": { "litellm_provider": "snowflake", @@ -30587,13 +30618,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/snowflake-llama-3.3-70b": { + "snowflake/snowflake-llama-3.3-70b": { + "max_tokens": 16384, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.00000072, + "output_cost_per_token": 0.00000072, "litellm_provider": "snowflake", - "max_input_tokens": 8000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat" - }, + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, "stability/sd3": { "litellm_provider": "stability", "mode": "image_generation", @@ -41223,6 +41258,44 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock_mantle/openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "cache_read_input_token_cost": 5.5e-07, + "output_cost_per_token": 3.3e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": ["/v1/responses"], + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["text"], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "cache_read_input_token_cost": 2.75e-07, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": ["/v1/responses"], + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["text"], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, @@ -41501,5 +41574,180 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_pdf_input": true - } -} + }, + "snowflake/claude-sonnet-4-5": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.0000003, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-sonnet-4-6": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.0000003, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-4-sonnet": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.0000003, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-4-opus": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000005, + "output_cost_per_token": 0.000025, + "cache_read_input_token_cost": 0.0000005, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/claude-haiku-4-5": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000005, + "cache_read_input_token_cost": 0.0000001, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-3-7-sonnet": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_read_input_token_cost": 0.0000003, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-4.1": { + "max_tokens": 16384, + "max_input_tokens": 300000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.000002, + "output_cost_per_token": 0.000008, + "cache_read_input_token_cost": 0.0000005, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5": { + "max_tokens": 16384, + "max_input_tokens": 300000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.00000125, + "output_cost_per_token": 0.00001, + "cache_read_input_token_cost": 0.000000125, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5-mini": { + "max_tokens": 16384, + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.0000003, + "output_cost_per_token": 0.0000012, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5-nano": { + "max_tokens": 16384, + "max_input_tokens": 5000000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.00000015, + "output_cost_per_token": 0.0000006, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/llama4-maverick": { + "max_tokens": 16384, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "input_cost_per_token": 0.00000024, + "output_cost_per_token": 0.00000097, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "snowflake/snowflake-arctic-embed-l-v2.0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "input_cost_per_token": 0.00000007, + "output_cost_per_token": 0.0, + "litellm_provider": "snowflake", + "mode": "embedding" + }, + "snowflake/snowflake-arctic-embed-m-v2.0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "input_cost_per_token": 0.00000007, + "output_cost_per_token": 0.0, + "litellm_provider": "snowflake", + "mode": "embedding" + } + } + diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py index 47b95c27ab7..c4f15ac4c3e 100644 --- a/tests/llm_translation/test_fireworks_ai_translation.py +++ b/tests/llm_translation/test_fireworks_ai_translation.py @@ -43,12 +43,14 @@ def test_map_openai_params_tool_choice(): def test_map_response_format(): """ - Test that the response format is translated correctly. + json_schema response_format is passed through to Fireworks unchanged. - h/t to https://github.com/DaveDeCaprio (@DaveDeCaprio) for the test case + Fireworks accepts the OpenAI strict json_schema shape natively. The earlier + downgrade to {type: json_object, schema: ...} silently dropped `strict` and + `name`, producing a request that Fireworks treats as "any valid JSON" per + its docs, disabling grammar-guided decoding. - Relevant Issue: https://github.com/BerriAI/litellm/issues/6797 - Fireworks AI Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting#step-1-import-libraries + Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting """ response_format = { "type": "json_schema", @@ -65,16 +67,7 @@ def test_map_response_format(): result = fireworks.map_openai_params( {"response_format": response_format}, {}, "some_model", drop_params=False ) - assert result == { - "response_format": { - "type": "json_object", - "schema": { - "properties": {"result": {"type": "boolean"}}, - "required": ["result"], - "type": "object", - }, - } - } + assert result == {"response_format": response_format} class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest): diff --git a/tests/local_testing/test_custom_llm.py b/tests/local_testing/test_custom_llm.py index 34ab6c043b9..ea15c3db9d0 100644 --- a/tests/local_testing/test_custom_llm.py +++ b/tests/local_testing/test_custom_llm.py @@ -44,7 +44,14 @@ from litellm import ( image_generation, ) from litellm.utils import ModelResponseIterator -from litellm.types.utils import ImageResponse, ImageObject, EmbeddingResponse +from litellm.types.utils import ( + ImageResponse, + ImageObject, + EmbeddingResponse, + ModelResponseStream, + StreamingChoices, + Delta, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -644,3 +651,82 @@ async def test_simple_aembedding(): "embedding": [0.1, 0.2, 0.3], "index": 1, } + + +# ── Tests for ModelResponseStream passthrough in custom providers (issue #27389) ── + + +class ModelResponseStreamLLM(MyCustomLLM): + """Subclass that overrides streaming/astreaming to yield ModelResponseStream directly.""" + + def __init__(self, finish_reason: str = "stop"): + self._finish_reason = finish_reason + + def streaming(self, *args, **kwargs) -> Iterator[ModelResponseStream]: # type: ignore + yield ModelResponseStream( + id="test-stream-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello world"), + finish_reason=self._finish_reason, + ) + ], + ) + + async def astreaming(self, *args, **kwargs) -> AsyncIterator[ModelResponseStream]: # type: ignore + yield ModelResponseStream( + id="test-stream-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello world"), + finish_reason=self._finish_reason, + ) + ], + ) + + +@pytest.mark.parametrize( + "finish_reason", ["stop", "tool_calls", "length", "content_filter"] +) +def test_custom_llm_streaming_model_response_stream(finish_reason): + my_custom_llm = ModelResponseStreamLLM(finish_reason=finish_reason) + litellm.custom_provider_map = [ + {"provider": "custom_llm", "custom_handler": my_custom_llm} + ] + resp = completion( + model="custom_llm/my-fake-model", + messages=[{"role": "user", "content": "Hello world!"}], + stream=True, + ) + + for chunk in resp: + print(chunk) + if chunk.choices[0].finish_reason is None: + assert isinstance(chunk.choices[0].delta.content, str) + else: + assert chunk.choices[0].finish_reason == finish_reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "finish_reason", ["stop", "tool_calls", "length", "content_filter"] +) +async def test_custom_llm_astreaming_model_response_stream(finish_reason): + my_custom_llm = ModelResponseStreamLLM(finish_reason=finish_reason) + litellm.custom_provider_map = [ + {"provider": "custom_llm", "custom_handler": my_custom_llm} + ] + resp = await litellm.acompletion( + model="custom_llm/my-fake-model", + messages=[{"role": "user", "content": "Hello world!"}], + stream=True, + ) + + async for chunk in resp: + print(chunk) + if chunk.choices[0].finish_reason is None: + assert isinstance(chunk.choices[0].delta.content, str) + else: + assert chunk.choices[0].finish_reason == finish_reason diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 46f2f89b3ce..1c041be0949 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -131,6 +131,7 @@ def test_default_api_base(): from litellm.litellm_core_utils.get_llm_provider_logic import ( _get_openai_compatible_provider_info, ) + from litellm.types.utils import LlmProviders # Patch environment variable to remove API base if it's set with patch.dict(os.environ, {}, clear=True): @@ -150,13 +151,13 @@ def test_default_api_base(): if api_base is None: continue - for other_provider in litellm.provider_list: - if other_provider != provider and provider != "{}_chat".format( + for other_provider in LlmProviders: + if other_provider.value != provider and provider != "{}_chat".format( other_provider.value ): - if provider == "codestral" and other_provider == "mistral": + if provider == "codestral" and other_provider.value == "mistral": continue - elif provider == "github" and other_provider == "azure": + elif provider == "github" and other_provider.value == "azure": continue assert other_provider.value not in api_base.replace("/openai", "") diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py new file mode 100644 index 00000000000..3c280c6ba92 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py @@ -0,0 +1,134 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, +) + +BEDROCK_REAL_MODEL = "eu.anthropic.claude-haiku-4-5-20251001-v1:0" +BEDROCK_LABEL = "claude-haiku-4-5" + + +def test_base_model_label_does_not_strip_bedrock_tools(): + """Regression for #29618. + + A Bedrock deployment whose ``model_info.base_model`` is a friendly label + (``claude-haiku-4-5``) must still advertise ``tools``/``tool_choice``. The label + on its own resolves to no tool support, so before the fix it stripped the + capability the real model id exposes, silently dropping function calling under + ``drop_params``.""" + params = get_supported_openai_params( + model=BEDROCK_REAL_MODEL, + custom_llm_provider="bedrock", + base_model=BEDROCK_LABEL, + ) + + assert params is not None + assert "tools" in params + assert "tool_choice" in params + + +def test_base_model_label_alone_lacks_bedrock_tools(): + """The label by itself does not advertise tools; this is what made the union + necessary. Guards against the discrepancy disappearing (and the regression test + above silently passing for the wrong reason).""" + params = get_supported_openai_params( + model=BEDROCK_LABEL, custom_llm_provider="bedrock" + ) + + assert params is not None + assert "tools" not in params + + +def test_base_model_is_additive_not_replacement(): + """``base_model`` may only add capabilities, never remove ones the real model has. + + Bedrock: real id supports ``tools`` but not the label's reasoning hint; the union + must contain the real model's ``tools`` regardless of the label being a subset.""" + real_only = set( + get_supported_openai_params( + model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock" + ) + ) + label_only = set( + get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock") + ) + combined = set( + get_supported_openai_params( + model=BEDROCK_REAL_MODEL, + custom_llm_provider="bedrock", + base_model=BEDROCK_LABEL, + ) + ) + + assert combined == real_only | label_only + assert real_only - label_only # the label really is a strict subset here + assert real_only <= combined + + +def test_base_model_adds_capabilities_the_real_model_lacks(): + """Regression for #27717 (the behavior the union must preserve). + + ``gemini-3.1-pro`` isn't in the cost map so it advertises no reasoning support, + but the registered ``gemini-3.1-pro-preview`` base_model does. The hint must add + ``reasoning_effort``/``thinking`` without the call erroring.""" + real_only = set( + get_supported_openai_params( + model="gemini-3.1-pro", custom_llm_provider="gemini" + ) + ) + assert "reasoning_effort" not in real_only + + combined = set( + get_supported_openai_params( + model="gemini-3.1-pro", + custom_llm_provider="gemini", + base_model="gemini-3.1-pro-preview", + ) + ) + assert "reasoning_effort" in combined + assert "thinking" in combined + + +def test_no_base_model_is_unchanged(): + """Omitting ``base_model`` must resolve purely from ``model``.""" + with_none = get_supported_openai_params( + model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None + ) + plain = get_supported_openai_params( + model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock" + ) + + assert with_none == plain + + +def test_base_model_equal_to_model_is_unchanged(): + """A ``base_model`` identical to ``model`` must not double-resolve or reorder.""" + plain = get_supported_openai_params( + model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock" + ) + same = get_supported_openai_params( + model=BEDROCK_REAL_MODEL, + custom_llm_provider="bedrock", + base_model=BEDROCK_REAL_MODEL, + ) + + assert same == plain + + +def test_azure_base_model_detection_preserved(): + """Azure relies on ``base_model`` for model-type detection when the deployment + name is opaque; the union must keep advertising the gpt-5 capabilities.""" + params = get_supported_openai_params( + model="my-opaque-deployment", + custom_llm_provider="azure", + base_model="azure/gpt-5.2", + ) + + assert params is not None + assert "reasoning_effort" in params + assert "tools" in params diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 63e2cb7f35c..b2002f9a0f9 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2118,3 +2118,172 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum(): f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. " "STOP enum was not normalised through map_finish_reason()." ) + + +@pytest.mark.parametrize( + "finish_reason", ["stop", "tool_calls", "length", "content_filter"] +) +def test_chunk_creator_passes_through_model_response_stream( + initialized_custom_stream_wrapper: CustomStreamWrapper, + finish_reason: str, +): + """ + chunk_creator must pass ModelResponseStream chunks from custom providers + straight through and preserve finish_reason exactly — not force-cast to GChunk. + Regression test for issue #27389. + """ + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider" + litellm._custom_providers.append("my-custom-provider") + + chunk = ModelResponseStream( + id="test-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello", role="assistant"), + finish_reason=finish_reason, + ) + ], + ) + + result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk) + + litellm._custom_providers.remove("my-custom-provider") + + assert result is not None + assert initialized_custom_stream_wrapper.received_finish_reason == finish_reason + + +def test_chunk_creator_drops_empty_finish_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ + A ModelResponseStream chunk with finish_reason but no content should return + None so finish_reason_handler() synthesises the final chunk — mirrors GChunk + behaviour via is_chunk_non_empty. + """ + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider" + litellm._custom_providers.append("my-custom-provider") + + chunk = ModelResponseStream( + id="test-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=""), + finish_reason="stop", + ) + ], + ) + + result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk) + + litellm._custom_providers.remove("my-custom-provider") + + assert result is None + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_chunk_creator_stops_iteration_on_trailing_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ + After received_finish_reason is set, any empty trailing chunk (e.g. provider + metadata flush) must raise StopIteration to end the stream cleanly. + """ + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider" + initialized_custom_stream_wrapper.received_finish_reason = "stop" + litellm._custom_providers.append("my-custom-provider") + + trailing_chunk = ModelResponseStream( + id="test-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=None), + finish_reason="stop", + ) + ], + ) + + with pytest.raises(StopIteration): + initialized_custom_stream_wrapper.chunk_creator(chunk=trailing_chunk) + + litellm._custom_providers.remove("my-custom-provider") + + +def test_chunk_creator_strips_finish_reason_from_content_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ + When content and finish_reason arrive in the same chunk, finish_reason must be + stripped so finish_reason_handler() emits it on the synthetic terminal chunk — + preventing two terminal chunks (double finish_reason bug). + """ + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider" + litellm._custom_providers.append("my-custom-provider") + + chunk = ModelResponseStream( + id="test-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello"), + finish_reason="stop", + ) + ], + ) + + result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk) + + litellm._custom_providers.remove("my-custom-provider") + + assert result is not None + assert ( + result.choices[0].finish_reason is None + ), "finish_reason must be stripped from content chunks to avoid double terminal chunks" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_chunk_creator_tool_calls_not_dropped_on_finish( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ + A terminal chunk with finish_reason="tool_calls" and delta.tool_calls must NOT + be silently dropped — tool_calls counts as content so the chunk is passed through + (with finish_reason stripped) rather than returning None. + """ + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider" + litellm._custom_providers.append("my-custom-provider") + + chunk = ModelResponseStream( + id="test-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_abc", + function=Function(name="get_weather", arguments='{"city":"NYC"}'), + type="function", + index=0, + ) + ], + ), + finish_reason="tool_calls", + ) + ], + ) + + result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk) + + litellm._custom_providers.remove("my-custom-provider") + + assert result is not None, "tool_calls chunk must not be dropped" + assert result.choices[0].delta.tool_calls is not None + assert result.choices[0].finish_reason is None + assert initialized_custom_stream_wrapper.received_finish_reason == "tool_calls" diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index 9d2787f78a7..59472d1a49d 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py @@ -1,5 +1,7 @@ +import urllib.parse from unittest.mock import patch +import litellm from litellm.llms.azure.image_edit.transformation import AzureImageEditConfig from litellm.types.router import GenericLiteLLMParams @@ -138,3 +140,96 @@ def test_azure_finalize_image_edit_strips_model_after_openai_transform(): assert data_out.get("prompt") == prompt assert data_out.get("n") == 1 assert len(files) >= 1 + + +# --------------------------------------------------------------------------- +# api_version fallback chain +# +# Pin the resolution order used by ``AzureImageEditConfig.get_complete_url``: +# litellm_params["api_version"] +# > litellm.api_version (module-global) +# > AZURE_API_VERSION env var +# > litellm.AZURE_DEFAULT_API_VERSION +# +# Before this fallback chain existed, image edit only read ``litellm_params`` +# and produced an unversioned URL when callers set api_version via the global +# or the env var (Azure then 404s with "Resource not found"). The chat path +# in ``litellm/llms/azure/common_utils.py`` already had this fallback. +# --------------------------------------------------------------------------- + + +_FALLBACK_API_BASE = "https://x.openai.azure.com" +_FALLBACK_MODEL = "gpt-image-1" + + +def _query_params(url: str) -> dict: + return dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query)) + + +def test_api_version_uses_litellm_params_first(monkeypatch): + monkeypatch.setattr(litellm, "api_version", "from-global", raising=False) + monkeypatch.setenv("AZURE_API_VERSION", "from-env") + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={"api_version": "from-params"}, + ) + + assert _query_params(url) == {"api-version": "from-params"} + + +def test_api_version_falls_back_to_litellm_global(monkeypatch): + monkeypatch.setattr(litellm, "api_version", "from-global", raising=False) + monkeypatch.setenv("AZURE_API_VERSION", "from-env") + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={}, + ) + + assert _query_params(url) == {"api-version": "from-global"} + + +def test_api_version_falls_back_to_env_var(monkeypatch): + monkeypatch.setattr(litellm, "api_version", None, raising=False) + monkeypatch.setenv("AZURE_API_VERSION", "from-env") + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={}, + ) + + assert _query_params(url) == {"api-version": "from-env"} + + +def test_api_version_falls_back_to_azure_default(monkeypatch): + monkeypatch.setattr(litellm, "api_version", None, raising=False) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={}, + ) + + assert _query_params(url) == {"api-version": litellm.AZURE_DEFAULT_API_VERSION} + + +def test_api_version_in_api_base_query_is_preserved(monkeypatch): + """``api_base`` already carrying ``?api-version=...`` must not be overridden.""" + monkeypatch.setattr(litellm, "api_version", None, raising=False) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=( + f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}" + "/images/edits?api-version=2024-05-01-preview" + ), + litellm_params={"api_version": "would-be-overridden"}, + ) + + assert _query_params(url) == {"api-version": "2024-05-01-preview"} diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py new file mode 100644 index 00000000000..e2133d56f89 --- /dev/null +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -0,0 +1,283 @@ +""" +Unit tests for Amazon Bedrock Mantle Responses API configuration. + +Mantle's gpt-5.5 / gpt-5.4 are served ONLY on the non-standard +`/openai/v1/responses` path. These tests lock the URL construction and +Bearer auth that make that routing work. +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import pytest + +import litellm +from litellm.llms.bedrock_mantle.responses.transformation import ( + BedrockMantleResponsesAPIConfig, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + + +class TestBedrockMantleResponsesURL: + def test_url_uses_region_from_env(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + + def test_url_normalizes_v1_suffix(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + assert "/v1/openai/v1/responses" not in url + url_trailing = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/v1/", + litellm_params={}, + ) + assert ( + url_trailing + == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + + def test_url_does_not_double_openai_v1(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + + def test_url_full_endpoint_base_not_doubled(self, monkeypatch): + # AWS model card tells users to set OPENAI_BASE_URL to the full endpoint. + # If copied into api_base, it must not be doubled. + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + assert url.count("/responses") == 1 + + def test_url_region_fallback_to_aws_region(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.setenv("AWS_REGION", "us-west-2") + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-west-2.api.aws/openai/v1/responses" + + def test_url_region_default_us_east_1(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses" + + +class TestBedrockMantleResponsesAuth: + def test_config_api_key_takes_priority(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key") + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams(api_key="config-key"), + ) + assert headers["Authorization"] == "Bearer config-key" + + def test_env_key_fallback(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key") + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer env-key" + + def test_bedrock_bearer_token_fallback(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bearer-key") + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer bearer-key" + + def test_missing_key_raises(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + with pytest.raises(ValueError, match="Bedrock Mantle API key"): + cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams(), + ) + + def test_custom_llm_provider(self): + cfg = BedrockMantleResponsesAPIConfig() + assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE + + def test_native_websocket_disabled(self): + # Mantle Responses has no realtime/websocket transport, so the config + # must opt out; otherwise realtime routing would try a socket Mantle + # does not serve. + cfg = BedrockMantleResponsesAPIConfig() + assert cfg.supports_native_websocket() is False + + def test_file_search_routes_to_emulation(self): + # Mantle cannot reach OpenAI's vector stores, so a native file_search + # tool forwarded as-is gets a 400. The config must opt out of native + # file_search so LiteLLM's emulation handles it instead of forwarding. + from litellm.responses.file_search.emulated_handler import ( + should_use_emulated_file_search, + ) + + cfg = BedrockMantleResponsesAPIConfig() + assert cfg.supports_native_file_search() is False + assert ( + should_use_emulated_file_search( + tools=[{"type": "file_search", "vector_store_ids": ["vs_1"]}], + provider_config=cfg, + ) + is True + ) + + +class TestBedrockMantleResponsesRegistry: + def test_registry_returns_config_for_gpt_5_5(self): + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="openai.gpt-5.5", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + + def test_registry_returns_config_for_gpt_5_4_enum(self): + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider=LlmProviders.BEDROCK_MANTLE, + model="openai.gpt-5.4", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + + def test_registry_returns_none_for_gpt_oss(self): + # Regression guard: gpt-oss must NOT get the native Responses config; it + # keeps the chat-completions emulation path (responses/main.py ~line 1109). + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="openai.gpt-oss-120b", + ) + assert cfg is None + + def test_registry_returns_none_for_gpt_oss_safeguard(self): + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="openai.gpt-oss-safeguard-20b", + ) + assert cfg is None + + def test_registry_returns_config_for_future_frontier_model(self): + # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6) must + # get the native Responses config without a code change. The gate allow-lists + # the openai.gpt- family (minus gpt-oss), so gpt-6 matches automatically. + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="openai.gpt-6", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + + @pytest.mark.parametrize( + "model", + [ + "nvidia.nemotron-nano-9b-v2", + "mistral.ministral-3-3b-instruct", + "google.gemma-3-27b-it", + "zai.glm-4.6", + ], + ) + def test_registry_returns_none_for_non_openai_models(self, model): + # Regression for the chat-only families on Mantle. These models 400 on + # /openai/v1/responses and are served on /v1/chat/completions, so the + # registry must NOT hand them the Responses config; they fall through to + # None and keep the chat-completions emulation. + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model=model, + ) + assert cfg is None + + def test_registry_returns_none_when_model_is_none(self): + # By-id operations (delete/get/cancel) call with model=None; keep returning + # None so those paths are unchanged. + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model=None, + ) + assert cfg is None + + +@pytest.fixture +def local_cost_map(monkeypatch): + """Force the bundled backup cost map and re-derive the provider model sets. + + ``litellm.model_cost`` is populated once at import time (here, from the + network-fetched ``main`` copy, which lags this branch). ``add_known_models`` + only re-buckets whatever is already in ``model_cost``, so the cost map must + first be reloaded from the local backup before the new keys appear. + """ + 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() + litellm.add_known_models() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +class TestBedrockMantleResponsesPricing: + def test_gpt_5_5_pricing_and_mode(self, local_cost_map): + info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") + assert info["mode"] == "responses" + assert info["input_cost_per_token"] == pytest.approx(5.5e-06) + assert info["output_cost_per_token"] == pytest.approx(3.3e-05) + assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07) + assert info["max_input_tokens"] == 272000 + + def test_gpt_5_4_pricing_and_mode(self, local_cost_map): + info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4") + assert info["mode"] == "responses" + assert info["input_cost_per_token"] == pytest.approx(2.75e-06) + assert info["output_cost_per_token"] == pytest.approx(1.65e-05) + assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07) + assert info["max_input_tokens"] == 272000 + + def test_models_registered(self, local_cost_map): + assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models + assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 2061522feff..ca340b5f275 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -496,3 +496,59 @@ def test_transform_tools_skips_non_function_tools(): "type": "object", "properties": {"id": {"type": "string"}}, } + + +def test_map_response_format_passes_json_schema_through_unchanged(): + """ + json_schema response_format must reach Fireworks unchanged. + + Regression guard for the prior downgrade to {type: json_object, schema: ...} + which silently dropped `strict` and `name` and disabled grammar-guided + decoding on the Fireworks side. + """ + config = FireworksAIConfig() + response_format = { + "type": "json_schema", + "json_schema": { + "name": "priority_classification", + "strict": True, + "schema": { + "type": "object", + "properties": { + "priority": { + "type": "string", + "enum": ["high", "medium", "low"], + } + }, + "required": ["priority"], + "additionalProperties": False, + }, + }, + } + + result = config.map_openai_params( + {"response_format": response_format}, + {}, + "fireworks_ai/accounts/fireworks/models/qwen3-32b", + drop_params=False, + ) + + rf = result["response_format"] + assert rf["type"] == "json_schema" + assert rf["json_schema"]["name"] == "priority_classification" + assert rf["json_schema"]["strict"] is True + assert rf["json_schema"]["schema"] == response_format["json_schema"]["schema"] + + +def test_map_response_format_json_object_unchanged(): + """ + The plain json_object form keeps working as before. + """ + config = FireworksAIConfig() + result = config.map_openai_params( + {"response_format": {"type": "json_object"}}, + {}, + "fireworks_ai/accounts/fireworks/models/qwen3-32b", + drop_params=False, + ) + assert result == {"response_format": {"type": "json_object"}} diff --git a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py b/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py index 4cf2429d737..6f215deed4e 100644 --- a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py +++ b/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py @@ -2,6 +2,7 @@ Tests for Gemini (Veo) video generation transformation. """ +import io import json import os from unittest.mock import MagicMock, Mock, patch @@ -132,6 +133,87 @@ class TestGeminiVideoConfig: assert data["parameters"]["durationSeconds"] == 8 assert data["parameters"]["resolution"] == "1080p" + def test_transform_video_create_request_image_goes_to_instance(self): + """Image belongs in instances[0], not in parameters (per Veo API).""" + prompt = "Animate this still" + api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning" + image_dict = {"bytesBase64Encoded": "aGVsbG8=", "mimeType": "image/jpeg"} + + data, _, _ = self.config.transform_video_create_request( + model="veo-3.0-generate-preview", + prompt=prompt, + api_base=api_base, + video_create_optional_request_params={ + "image": image_dict, + "aspectRatio": "16:9", + "durationSeconds": 4, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["instances"][0]["prompt"] == prompt + assert data["instances"][0]["image"] == image_dict + assert "image" not in data.get("parameters", {}) + assert data["parameters"]["aspectRatio"] == "16:9" + assert data["parameters"]["durationSeconds"] == 4 + + def test_transform_video_create_request_image_filelike_goes_to_instance(self): + """File-like image (BytesIO) gets base64-encoded into instances[0]['image'].""" + prompt = "Animate this still" + api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning" + # 1x1 PNG (8 bytes after magic + minimal IHDR is not legal — but the + # transformer only cares that ImageEditRequestUtils can sniff a MIME and + # that .read() returns bytes; an explicit name="image.jpeg" hands the + # MIME sniffer a clean answer regardless of payload). + image_bytes = b"\xff\xd8\xff\xe0fake-jpeg-bytes" + image_file = io.BytesIO(image_bytes) + image_file.name = "still.jpeg" + + data, _, _ = self.config.transform_video_create_request( + model="veo-3.0-generate-preview", + prompt=prompt, + api_base=api_base, + video_create_optional_request_params={ + "image": image_file, + "aspectRatio": "16:9", + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + # File-like took the _convert_image_to_gemini_format branch and landed + # in instances[0]["image"], not in parameters. + instance_image = data["instances"][0]["image"] + assert isinstance(instance_image, dict) + assert instance_image["mimeType"].startswith("image/") + assert instance_image["bytesBase64Encoded"] + # Round-trip the base64 — should equal the original bytes. + import base64 + + assert base64.b64decode(instance_image["bytesBase64Encoded"]) == image_bytes + assert "image" not in data.get("parameters", {}) + + def test_transform_video_create_request_image_none_is_dropped(self): + """Explicit image=None is popped and never reaches parameters.""" + prompt = "no image at all" + api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning" + + data, _, _ = self.config.transform_video_create_request( + model="veo-3.0-generate-preview", + prompt=prompt, + api_base=api_base, + video_create_optional_request_params={ + "image": None, + "aspectRatio": "16:9", + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "image" not in data["instances"][0] + assert "image" not in data.get("parameters", {}) + def test_map_openai_params(self): """Test parameter mapping from OpenAI format to Veo format.""" openai_params = { diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 560796ea58d..8a072fa5097 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -104,6 +104,23 @@ class TestHuggingFaceEmbedding: assert "source_sentence" not in str(request_data) assert "sentences" not in str(request_data) + def test_embedding_allows_special_token_looking_input(self): + input_text = ["hello <|fim_prefix|> world"] + + response = litellm.embedding( + model=self.model, + input=input_text, + input_type="embed", + ) + + self.mock_http.assert_called_once() + post_call_args = self.mock_http.call_args + request_data = json.loads(post_call_args[1]["data"]) + + assert request_data["inputs"] == input_text + assert response.usage.prompt_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + def test_embedding_with_sentence_similarity_task(self): """Test embedding when task type is sentence-similarity (requires 2+ sentences)""" diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 6f32c4ca340..74888e6cd9e 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -201,9 +201,12 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache == "existing_cache_name" - # Verify cache key was generated with tools and model + # Verify cache key was generated with tools, tool_choice and model mock_cache_obj.get_cache_key.assert_called_once_with( - messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro" + messages=cached_messages, + tools=self.sample_tools, + tool_choice=None, + model="gemini-1.5-pro", ) @pytest.mark.parametrize( @@ -474,9 +477,12 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache == "existing_cache_name" - # Verify cache key was generated with tools and model + # Verify cache key was generated with tools, tool_choice and model mock_cache_obj.get_cache_key.assert_called_once_with( - messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro" + messages=cached_messages, + tools=self.sample_tools, + tool_choice=None, + model="gemini-1.5-pro", ) @pytest.mark.asyncio @@ -800,6 +806,546 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + def test_check_and_create_cache_tool_choice_popped_from_optional_params( + self, custom_llm_provider + ): + """tool_choice is popped from optional_params when cached messages exist.""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} + + with patch.object( + self.context_caching, "check_cache", return_value="existing_cache" + ): + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert "tool_choice" not in optional_params + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( + self, custom_llm_provider + ): + """tool_choice is NOT popped when there are no cached messages (early return).""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + mock_separate.return_value = ([], self.sample_messages) + + tool_choice = {"functionCallingConfig": {"mode": "AUTO"}} + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = tool_choice + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert optional_params.get("tool_choice") == tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params( + self, custom_llm_provider + ): + """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params.""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} + + with patch.object( + self.context_caching, "async_check_cache", return_value="existing_cache" + ): + await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert "tool_choice" not in optional_params + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( + self, custom_llm_provider + ): + """Async equivalent of test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages.""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + mock_separate.return_value = ([], self.sample_messages) + + tool_choice = {"functionCallingConfig": {"mode": "AUTO"}} + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = tool_choice + + await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert optional_params.get("tool_choice") == tool_choice + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_and_create_cache_tool_choice_in_request_body( + self, + mock_get_token_url, + mock_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """End-to-end: tool_choice ends up as `toolConfig` on the cache-creation HTTP POST body.""" + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None # cache miss -> create new + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_client.post.return_value = mock_response + + tool_choice = {"functionCallingConfig": {"mode": "ANY"}} + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = tool_choice + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + self.mock_client.post.assert_called_once() + call_args = self.mock_client.post.call_args + assert call_args.kwargs["json"]["tools"] == self.sample_tools + assert call_args.kwargs["json"]["toolConfig"] == tool_choice + mock_cache_obj.get_cache_key.assert_called_once_with( + messages=cached_messages, + tools=self.sample_tools, + tool_choice=tool_choice, + model="gemini-1.5-pro", + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "async_check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_and_create_cache_tool_choice_in_request_body( + self, + mock_get_token_url, + mock_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """Async equivalent of test_check_and_create_cache_tool_choice_in_request_body.""" + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_async_client.post = AsyncMock(return_value=mock_response) + + tool_choice = {"functionCallingConfig": {"mode": "ANY"}} + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = tool_choice + + await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + call_args = self.mock_async_client.post.call_args + assert call_args.kwargs["json"]["tools"] == self.sample_tools + assert call_args.kwargs["json"]["toolConfig"] == tool_choice + mock_cache_obj.get_cache_key.assert_called_once_with( + messages=cached_messages, + tools=self.sample_tools, + tool_choice=tool_choice, + model="gemini-1.5-pro", + ) + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_and_create_cache_omits_tool_config_when_tool_choice_unset( + self, + mock_get_token_url, + mock_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """When the caller didn't pass tool_choice, toolConfig must NOT appear in the cache body.""" + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_client.post.return_value = mock_response + + optional_params = self.sample_optional_params.copy() + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + call_args = self.mock_client.post.call_args + assert "tools" in call_args.kwargs["json"] + assert "toolConfig" not in call_args.kwargs["json"] + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_and_create_cache_tool_choice_function_pin( + self, + mock_get_token_url, + mock_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """tool_choice as a function-pin dict survives the cache body intact.""" + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_client.post.return_value = mock_response + + function_pin = { + "functionCallingConfig": { + "mode": "ANY", + "allowed_function_names": ["get_current_weather"], + } + } + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = function_pin + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + call_args = self.mock_client.post.call_args + assert call_args.kwargs["json"]["toolConfig"] == function_pin + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_and_create_cache_tool_choice_typed_constructor( + self, + mock_get_token_url, + mock_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """Exercise the actual ToolConfig(FunctionCallingConfig(...)) constructor that map_tool_choice_values produces. + + ToolConfig / FunctionCallingConfig are TypedDicts (litellm/types/llms/vertex_ai.py:158, 277) + so this is functionally identical to the dict-literal tests above at + runtime — but exercising the typed constructor pins the test to the + same call shape map_tool_choice_values uses and auto-follows if + either type ever migrates to a Pydantic model upstream. + """ + from litellm.types.llms.vertex_ai import ( + FunctionCallingConfig, + ToolConfig, + ) + + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_client.post.return_value = mock_response + + tool_choice = ToolConfig( + functionCallingConfig=FunctionCallingConfig(mode="ANY") + ) + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = tool_choice + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + call_args = self.mock_client.post.call_args + assert call_args.kwargs["json"]["toolConfig"] == tool_choice + assert call_args.kwargs["json"]["toolConfig"] == { + "functionCallingConfig": {"mode": "ANY"} + } + mock_cache_obj.get_cache_key.assert_called_once_with( + messages=cached_messages, + tools=self.sample_tools, + tool_choice=tool_choice, + model="gemini-1.5-pro", + ) + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @patch.object(ContextCachingEndpoints, "check_cache") + def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys( + self, + mock_check_cache, + mock_separate, + custom_llm_provider, + ): + """Two requests with different tool_choice values must produce different cache keys. + + Runs the real local_cache_obj.get_cache_key to verify the hashed + output actually differs — mocking it would only prove that distinct + arguments are forwarded, not that they produce distinct keys. + """ + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + mock_check_cache.return_value = "existing_cache" + + auto_tool_choice = {"functionCallingConfig": {"mode": "AUTO"}} + any_tool_choice = {"functionCallingConfig": {"mode": "ANY"}} + for choice in (auto_tool_choice, any_tool_choice): + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = choice + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + check_cache_calls = mock_check_cache.call_args_list + assert len(check_cache_calls) == 2 + first_cache_key = check_cache_calls[0].kwargs["cache_key"] + second_cache_key = check_cache_calls[1].kwargs["cache_key"] + assert first_cache_key != second_cache_key + @pytest.mark.parametrize( "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py new file mode 100644 index 00000000000..b93f0d56f8e --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py @@ -0,0 +1,211 @@ +""" +Tests for the MCP elicitation handler. + +Covers the gateway-mode relay logic (`elicitation/create` requests from an +upstream MCP server being forwarded to the connected downstream client) as +well as the decline paths used in tool-bridge mode or when the downstream +client lacks the requested elicitation capability. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from mcp.types import ( + ElicitRequestFormParams, + ElicitRequestURLParams, + ElicitResult, + ErrorData, +) + +from litellm.proxy._experimental.mcp_server import elicitation_handler +from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + _relay_elicitation_to_downstream, + handle_elicitation_request, +) + + +def _form_params(message: str = "fill the form") -> ElicitRequestFormParams: + return ElicitRequestFormParams( + mode="form", + message=message, + requestedSchema={"type": "object", "properties": {}}, + ) + + +def _url_params(message: str = "please authorize") -> ElicitRequestURLParams: + return ElicitRequestURLParams( + mode="url", + message=message, + url="https://example.com/oauth", + elicitationId="elc-1", + ) + + +def _caps(*, url=True, form=True) -> SimpleNamespace: + elicit = SimpleNamespace( + url=object() if url else None, + form=object() if form else None, + ) + return SimpleNamespace(elicitation=elicit) + + +class TestHandleElicitationRequest: + async def test_should_decline_when_no_downstream_session(self): + result = await handle_elicitation_request( + context=SimpleNamespace(), + params=_form_params(), + downstream_session=None, + ) + assert isinstance(result, ElicitResult) + assert result.action == "decline" + + async def test_should_relay_to_downstream_when_session_present(self): + accepted = ElicitResult(action="accept", content={"name": "ada"}) + session = SimpleNamespace(elicit_form=AsyncMock(return_value=accepted)) + + result = await handle_elicitation_request( + context=SimpleNamespace(), + params=_form_params(), + downstream_session=session, + downstream_capabilities=None, + ) + + assert result is accepted + session.elicit_form.assert_awaited_once() + + async def test_should_return_error_data_when_unavailable(self, monkeypatch): + monkeypatch.setattr(elicitation_handler, "MCP_ELICITATION_AVAILABLE", False) + result = await handle_elicitation_request( + context=SimpleNamespace(), + params=_form_params(), + downstream_session=SimpleNamespace(), + ) + assert isinstance(result, ErrorData) + assert "not available" in result.message + + async def test_should_return_error_data_on_unexpected_failure(self): + class _ExplodingParams: + mode = "form" + + @property + def message(self): + raise RuntimeError("boom") + + result = await handle_elicitation_request( + context=SimpleNamespace(), + params=_ExplodingParams(), + downstream_session=None, + ) + assert isinstance(result, ErrorData) + assert "boom" in result.message + + +class TestRelayElicitationToDownstream: + async def test_should_relay_form_mode(self): + accepted = ElicitResult(action="accept", content={"name": "ada"}) + session = SimpleNamespace(elicit_form=AsyncMock(return_value=accepted)) + + params = _form_params("collect name") + result = await _relay_elicitation_to_downstream( + params=params, + downstream_session=session, + downstream_capabilities=_caps(form=True), + ) + + assert result is accepted + session.elicit_form.assert_awaited_once() + _, kwargs = session.elicit_form.call_args + assert kwargs["message"] == "collect name" + assert kwargs["requestedSchema"] == params.requestedSchema + + async def test_should_relay_url_mode(self): + accepted = ElicitResult(action="accept") + session = SimpleNamespace(elicit_url=AsyncMock(return_value=accepted)) + + result = await _relay_elicitation_to_downstream( + params=_url_params(), + downstream_session=session, + downstream_capabilities=_caps(url=True), + ) + + assert result is accepted + session.elicit_url.assert_awaited_once() + _, kwargs = session.elicit_url.call_args + assert kwargs["url"] == "https://example.com/oauth" + assert kwargs["elicitation_id"] == "elc-1" + + async def test_should_use_generic_elicit_for_unknown_param_type(self): + accepted = ElicitResult(action="accept") + session = SimpleNamespace(elicit=AsyncMock(return_value=accepted)) + + # A bare params object that is neither Form nor URL params triggers + # the generic fallback path. + params = SimpleNamespace(mode="form", message="hi", requestedSchema={}) + result = await _relay_elicitation_to_downstream( + params=params, + downstream_session=session, + downstream_capabilities=None, + ) + + assert result is accepted + session.elicit.assert_awaited_once() + + async def test_should_decline_when_elicitation_unsupported(self): + session = SimpleNamespace(elicit_form=AsyncMock()) + caps = SimpleNamespace(elicitation=None) + + result = await _relay_elicitation_to_downstream( + params=_form_params(), + downstream_session=session, + downstream_capabilities=caps, + ) + + assert isinstance(result, ElicitResult) + assert result.action == "decline" + session.elicit_form.assert_not_awaited() + + async def test_should_decline_url_mode_when_url_unsupported(self): + session = SimpleNamespace(elicit_url=AsyncMock()) + + result = await _relay_elicitation_to_downstream( + params=_url_params(), + downstream_session=session, + downstream_capabilities=_caps(url=False, form=True), + ) + + assert isinstance(result, ElicitResult) + assert result.action == "decline" + session.elicit_url.assert_not_awaited() + + async def test_should_decline_form_mode_when_form_unsupported(self): + session = SimpleNamespace(elicit_form=AsyncMock()) + + result = await _relay_elicitation_to_downstream( + params=_form_params(), + downstream_session=session, + downstream_capabilities=_caps(url=True, form=False), + ) + + assert isinstance(result, ElicitResult) + assert result.action == "decline" + session.elicit_form.assert_not_awaited() + + async def test_should_decline_when_downstream_relay_raises(self): + session = SimpleNamespace( + elicit_form=AsyncMock(side_effect=RuntimeError("transport closed")) + ) + + result = await _relay_elicitation_to_downstream( + params=_form_params(), + downstream_session=session, + downstream_capabilities=_caps(form=True), + ) + + assert isinstance(result, ElicitResult) + assert result.action == "decline" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 04ff1e4be20..363948ff4e6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -16,7 +16,6 @@ from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager from litellm.proxy._types import UserAPIKeyAuth @@ -549,11 +548,7 @@ class TestHookHeaderMergePriority: captured_extra_headers: Dict[str, Any] = {} async def fake_create_mcp_client( - server, - mcp_auth_header=None, - extra_headers=None, - stdio_env=None, - subject_token=None, + server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs ): captured_extra_headers["value"] = extra_headers mock_client = MagicMock() @@ -593,11 +588,7 @@ class TestHookHeaderMergePriority: captured_extra_headers: Dict[str, Any] = {} async def fake_create_mcp_client( - server, - mcp_auth_header=None, - extra_headers=None, - stdio_env=None, - subject_token=None, + server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs ): captured_extra_headers["value"] = extra_headers mock_client = MagicMock() @@ -643,11 +634,7 @@ class TestHookHeaderMergePriority: captured_extra_headers: Dict[str, Any] = {} async def fake_create_mcp_client( - server, - mcp_auth_header=None, - extra_headers=None, - stdio_env=None, - subject_token=None, + server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs ): captured_extra_headers["value"] = extra_headers mock_client = MagicMock() @@ -703,11 +690,7 @@ class TestHookHeaderMergePriority: captured_extra_headers: Dict[str, Any] = {} async def fake_create_mcp_client( - server, - mcp_auth_header=None, - extra_headers=None, - stdio_env=None, - subject_token=None, + server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs ): captured_extra_headers["value"] = extra_headers mock_client = MagicMock() @@ -755,11 +738,7 @@ class TestHookHeaderMergePriority: captured_extra_headers: Dict[str, Any] = {} async def fake_create_mcp_client( - server, - mcp_auth_header=None, - extra_headers=None, - stdio_env=None, - subject_token=None, + server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs ): captured_extra_headers["value"] = extra_headers mock_client = MagicMock() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py new file mode 100644 index 00000000000..78aee7b534f --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py @@ -0,0 +1,254 @@ +""" +Tests for the MCP sampling completion pipeline. + +Covers building the internal `acompletion` kwargs from MCP request params +(messages, sampling options, tools, tool choice, metadata), routing the call +through the proxy router / guardrails, and the end-to-end +`handle_sampling_create_message` success and error-propagation behaviour. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from mcp.types import CreateMessageResult, ErrorData + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _build_completion_kwargs, + _run_guardrails_and_call_llm, + handle_sampling_create_message, +) + + +def _params(**overrides): + base = dict( + messages=[ + SimpleNamespace( + role="user", content=SimpleNamespace(type="text", text="hi") + ) + ], + systemPrompt="be concise", + maxTokens=128, + temperature=None, + stopSequences=None, + tools=None, + toolChoice=None, + metadata=None, + modelPreferences=None, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def _passthrough_add_data(): + async def _add(data, **kwargs): + return data + + return _add + + +class TestBuildCompletionKwargs: + async def test_should_include_sampling_options_and_tools(self): + params = _params( + temperature=0.3, + stopSequences=["STOP"], + tools=[ + SimpleNamespace( + name="search", description="d", inputSchema={"type": "object"} + ) + ], + toolChoice=SimpleNamespace(mode="required"), + metadata={"trace": "abc"}, + ) + with patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + side_effect=_passthrough_add_data(), + ): + kwargs = await _build_completion_kwargs( + params=params, + model="gpt-4o", + user_api_key_auth=SimpleNamespace(user_id="u1"), + raw_headers=None, + client_ip=None, + ) + + assert kwargs["model"] == "gpt-4o" + assert kwargs["max_tokens"] == 128 + assert kwargs["temperature"] == 0.3 + assert kwargs["stop"] == ["STOP"] + assert kwargs["tools"][0]["function"]["name"] == "search" + assert kwargs["tool_choice"] == "required" + assert kwargs["metadata"]["mcp_metadata"] == {"trace": "abc"} + assert kwargs["user"] == "u1" + assert kwargs["messages"][0] == {"role": "system", "content": "be concise"} + + async def test_should_omit_optional_fields_when_unset(self): + with patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + side_effect=_passthrough_add_data(), + ): + kwargs = await _build_completion_kwargs( + params=_params(), + model="gpt-4o", + user_api_key_auth=SimpleNamespace(user_id=None), + raw_headers=None, + client_ip=None, + ) + + assert "temperature" not in kwargs + assert "stop" not in kwargs + assert "tools" not in kwargs + assert "tool_choice" not in kwargs + assert kwargs["metadata"] == {} + + +class TestRunGuardrailsAndCallLlm: + async def test_should_route_through_llm_router_when_available(self): + router = MagicMock() + router.acompletion = AsyncMock(return_value="router-response") + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + patch("litellm.proxy.proxy_server.llm_router", router), + ): + result = await _run_guardrails_and_call_llm( + completion_kwargs={"model": "gpt-4o", "messages": []}, + user_api_key_auth=SimpleNamespace(), + ) + + assert result == "router-response" + router.acompletion.assert_awaited_once() + + async def test_should_propagate_guardrail_rejection(self): + plo = MagicMock() + plo.pre_call_hook = AsyncMock(side_effect=ValueError("blocked by guardrail")) + with patch("litellm.proxy.proxy_server.proxy_logging_obj", plo): + with pytest.raises(ValueError, match="blocked by guardrail"): + await _run_guardrails_and_call_llm( + completion_kwargs={"model": "gpt-4o", "messages": []}, + user_api_key_auth=SimpleNamespace(), + ) + + +class TestHandleSamplingCreateMessagePipeline: + async def test_should_return_message_result_on_success(self): + auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok") + response = SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace( + content="the answer is 42", tool_calls=None + ), + finish_reason="stop", + ) + ], + model="gpt-4o", + ) + with ( + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences", + return_value="gpt-4o", + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs", + new_callable=AsyncMock, + return_value={"model": "gpt-4o", "messages": []}, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._run_guardrails_and_call_llm", + new_callable=AsyncMock, + return_value=response, + ), + ): + result = await handle_sampling_create_message( + context=MagicMock(), + params=_params(), + default_model="gpt-4o", + user_api_key_auth=auth, + ) + + assert isinstance(result, CreateMessageResult) + assert result.content.text == "the answer is 42" + assert result.stopReason == "endTurn" + + async def test_should_reraise_known_proxy_exceptions(self): + from litellm.exceptions import RateLimitError + + auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok") + with ( + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences", + return_value="gpt-4o", + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs", + new_callable=AsyncMock, + side_effect=RateLimitError( + "rate limited", llm_provider="openai", model="gpt-4o" + ), + ), + ): + with pytest.raises(RateLimitError): + await handle_sampling_create_message( + context=MagicMock(), + params=_params(), + default_model="gpt-4o", + user_api_key_auth=auth, + ) + + async def test_should_return_error_data_on_unexpected_failure(self): + auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok") + with ( + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences", + return_value="gpt-4o", + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs", + new_callable=AsyncMock, + side_effect=RuntimeError("kaboom"), + ), + ): + result = await handle_sampling_create_message( + context=MagicMock(), + params=_params(), + default_model="gpt-4o", + user_api_key_auth=auth, + ) + + assert isinstance(result, ErrorData) + assert "kaboom" in result.message + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py new file mode 100644 index 00000000000..f141cb2e316 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py @@ -0,0 +1,327 @@ +""" +Tests for MCP sampling handler model-access enforcement. + +Verifies that handle_sampling_create_message and _check_model_access +enforce the same model-permission checks as regular /chat/completions +calls, preventing a malicious upstream MCP server from requesting +inference on models the caller's API key is not authorized to use. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _check_model_access, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_user_api_key_auth( + *, + models=None, + team_id=None, + team_model_aliases=None, + api_key="sk-test-key", + token=None, + user_role=None, +): + """Build a minimal UserAPIKeyAuth-like object for tests.""" + auth = MagicMock() + auth.models = models or [] + auth.team_id = team_id + auth.team_model_aliases = team_model_aliases or {} + auth.access_group_ids = [] + auth.api_key = api_key + auth.token = token + auth.user_role = user_role + return auth + + +# --------------------------------------------------------------------------- +# _check_model_access +# --------------------------------------------------------------------------- + + +class TestCheckModelAccess: + """Tests for the _check_model_access helper that gates sampling requests.""" + + @pytest.mark.asyncio + async def test_should_return_none_when_no_auth_context(self): + """No auth context means no restriction — pass through.""" + result = await _check_model_access("gpt-4o", user_api_key_auth=None) + assert result is None + + @pytest.mark.asyncio + async def test_should_allow_model_when_key_has_access(self): + """Key with explicit model access should be allowed.""" + auth = _make_user_api_key_auth(models=["gpt-4o", "gpt-3.5-turbo"]) + + with patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new_callable=AsyncMock, + return_value=True, + ) as mock_check: + result = await _check_model_access("gpt-4o", user_api_key_auth=auth) + + assert result is None + mock_check.assert_awaited_once() + + @pytest.mark.asyncio + async def test_should_deny_model_when_key_lacks_access(self): + """Key without model access should be denied with ErrorData.""" + from litellm.proxy._types import ProxyException + + auth = _make_user_api_key_auth(models=["gpt-3.5-turbo"]) + + with patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new_callable=AsyncMock, + side_effect=ProxyException( + message="key not allowed to access model", + type="key_model_access_denied", + param="model", + code=401, + ), + ): + result = await _check_model_access("gpt-4o", user_api_key_auth=auth) + + # Should return ErrorData, not raise + assert result is not None + assert result.code == -1 + assert "Model access denied" in result.message + assert "gpt-4o" in result.message + + @pytest.mark.asyncio + async def test_should_allow_wildcard_model_access(self): + """Key with wildcard model access should allow any model.""" + auth = _make_user_api_key_auth(models=["*"]) + + with patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new_callable=AsyncMock, + return_value=True, + ): + result = await _check_model_access( + "claude-3-opus-20240229", user_api_key_auth=auth + ) + + assert result is None + + @pytest.mark.asyncio + async def test_should_deny_expensive_model_requested_by_malicious_server(self): + """Simulates the attack: malicious MCP server hints at an expensive model + the caller's key is restricted from using.""" + from litellm.proxy._types import ProxyException + + # Key only has access to cheap models + auth = _make_user_api_key_auth(models=["gpt-3.5-turbo"]) + + with patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new_callable=AsyncMock, + side_effect=ProxyException( + message="key not allowed to access model. This key can only access models=['gpt-3.5-turbo']. Tried to access claude-3-opus-20240229", + type="key_model_access_denied", + param="model", + code=401, + ), + ): + result = await _check_model_access( + "claude-3-opus-20240229", user_api_key_auth=auth + ) + + assert result is not None + assert result.code == -1 + assert "claude-3-opus-20240229" in result.message + + @pytest.mark.asyncio + async def test_should_deny_empty_oauth_passthrough_placeholder(self): + """Regression: process_mcp_request() returns an empty UserAPIKeyAuth() + for OAuth2 upstream-token passthrough. The None check alone is not + sufficient — the empty placeholder is truthy but has no api_key, no + token, and an empty models list. can_key_call_model() would treat + that as all-model access, letting an OAuth-only user trigger sampling + calls on any proxy model without a LiteLLM key or budget.""" + # Simulate the empty placeholder from process_mcp_request() + auth = _make_user_api_key_auth( + models=[], + api_key=None, + token=None, + user_role=None, + ) + + result = await _check_model_access("gpt-4o", user_api_key_auth=auth) + + # Must be denied — not passed through to can_key_call_model + assert result is not None + assert result.code == -1 + assert "sampling requires a valid LiteLLM" in result.message + + @pytest.mark.asyncio + async def test_should_allow_proxy_admin_even_without_api_key(self): + """Proxy admins may not have a traditional api_key but should still + be allowed to use sampling.""" + auth = _make_user_api_key_auth( + models=[], + api_key=None, + token=None, + user_role="proxy_admin", + ) + + with patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new_callable=AsyncMock, + return_value=True, + ): + result = await _check_model_access("gpt-4o", user_api_key_auth=auth) + + assert result is None + + +# --------------------------------------------------------------------------- +# handle_sampling_create_message — auth + budget gating +# --------------------------------------------------------------------------- + + +class TestSamplingAuthAndBudgetGating: + + @pytest.mark.asyncio + async def test_should_deny_when_no_auth_context(self): + """Sampling must reject calls with no user_api_key_auth.""" + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + + params = MagicMock() + params.modelPreferences = None + params.messages = [] + params.systemPrompt = None + params.maxTokens = 100 + params.temperature = None + params.stopSequences = None + params.tools = None + params.toolChoice = None + params.metadata = None + + result = await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o", + user_api_key_auth=None, + ) + + assert result is not None + assert result.code == -1 + assert "authenticated" in result.message.lower() + + @pytest.mark.asyncio + async def test_should_run_budget_checks(self): + """Sampling must call _run_budget_checks after model access check.""" + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + + auth = _make_user_api_key_auth(models=["gpt-4o"]) + params = MagicMock() + params.modelPreferences = None + params.messages = [] + params.systemPrompt = None + params.maxTokens = 100 + params.temperature = None + params.stopSequences = None + params.tools = None + params.toolChoice = None + params.metadata = None + + with ( + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks", + new_callable=AsyncMock, + return_value=None, + ) as mock_budget, + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences", + return_value="gpt-4o", + ), + patch( + "litellm.proxy.proxy_server.llm_router", + new=None, + ), + patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=MagicMock( + choices=[ + MagicMock( + message=MagicMock(content="hi", tool_calls=None), + finish_reason="stop", + ) + ], + model="gpt-4o", + ), + ), + ): + await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o", + user_api_key_auth=auth, + ) + + mock_budget.assert_awaited_once() + + @pytest.mark.asyncio + async def test_should_deny_over_budget_caller(self): + """When _run_budget_checks returns ErrorData, sampling must return it.""" + from mcp.types import ErrorData + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + + auth = _make_user_api_key_auth(models=["gpt-4o"]) + params = MagicMock() + params.modelPreferences = None + params.messages = [] + params.systemPrompt = None + params.maxTokens = 100 + params.temperature = None + params.stopSequences = None + params.tools = None + params.toolChoice = None + params.metadata = None + + budget_error = ErrorData(code=-1, message="ExceededBudget: over limit") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks", + new_callable=AsyncMock, + return_value=budget_error, + ), + patch( + "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences", + return_value="gpt-4o", + ), + ): + result = await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o", + user_api_key_auth=auth, + ) + + assert result is budget_error + assert "ExceededBudget" in result.message diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py new file mode 100644 index 00000000000..0c8f7bd4814 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py @@ -0,0 +1,91 @@ +""" +Tests for MCP sampling model resolution (hint matching and fallback chain). + +`_resolve_model_from_preferences` first tries to match upstream model hints +against the proxy's available models (direct then substring), then priority +scoring, then the caller default, the first available model, and finally the +configured `default_mcp_sampling_model` before raising. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _resolve_model_from_preferences, +) + + +def _prefs(*, hints=None, cost=None, speed=None, intelligence=None): + return SimpleNamespace( + hints=hints or [], + costPriority=cost, + speedPriority=speed, + intelligencePriority=intelligence, + ) + + +class TestHintMatching: + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch("litellm.model_list", [{"model_name": "gpt-4o"}, {"model_name": "claude-3"}]) + def test_should_match_hint_as_substring(self): + prefs = _prefs(hints=[SimpleNamespace(name="gpt-4")]) + assert _resolve_model_from_preferences(prefs) == "gpt-4o" + + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch("litellm.model_list", ["gpt-4o", "claude-3"]) + def test_should_match_hint_against_string_model_list_entries(self): + prefs = _prefs(hints=[SimpleNamespace(name="claude-3")]) + assert _resolve_model_from_preferences(prefs) == "claude-3" + + @patch("litellm.model_list", None) + def test_should_use_router_model_names(self): + router = MagicMock() + router.get_model_names.return_value = ["router-gpt", "router-claude"] + with patch("litellm.proxy.proxy_server.llm_router", router): + prefs = _prefs(hints=[SimpleNamespace(name="router-claude")]) + assert _resolve_model_from_preferences(prefs) == "router-claude" + + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch("litellm.model_list", [{"model_name": "gpt-4o"}]) + def test_should_skip_hint_without_name(self): + prefs = _prefs(hints=[SimpleNamespace()]) # hint has no `.name` + assert ( + _resolve_model_from_preferences(prefs, default_model="gpt-4o") == "gpt-4o" + ) + + +class TestFallbackChain: + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch( + "litellm.model_list", [{"model_name": "first-model"}, {"model_name": "second"}] + ) + def test_should_fall_back_to_first_available_when_no_default(self): + prefs = _prefs(hints=[SimpleNamespace(name="no-such")]) + assert _resolve_model_from_preferences(prefs) == "first-model" + + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch("litellm.model_list", []) + def test_should_use_configured_default_sampling_model(self, monkeypatch): + import litellm + + monkeypatch.setattr( + litellm, "default_mcp_sampling_model", "fallback-model", raising=False + ) + prefs = _prefs() + assert _resolve_model_from_preferences(prefs) == "fallback-model" + + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch("litellm.model_list", []) + def test_should_raise_when_nothing_resolvable(self, monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "default_mcp_sampling_model", None, raising=False) + prefs = _prefs() + with pytest.raises(ValueError, match="No model could be resolved"): + _resolve_model_from_preferences(prefs) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py new file mode 100644 index 00000000000..24309ed0460 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py @@ -0,0 +1,248 @@ +""" +Tests for MCP sampling handler priority-based model selection. + +Verifies that _resolve_model_from_preferences honours costPriority, +speedPriority, and intelligencePriority when hints don't match, +per the MCP spec. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _has_priorities, + _resolve_model_from_preferences, + _select_model_by_priority, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _prefs(*, hints=None, cost=None, speed=None, intelligence=None): + """Build a minimal ModelPreferences-like object.""" + return SimpleNamespace( + hints=hints or [], + costPriority=cost, + speedPriority=speed, + intelligencePriority=intelligence, + ) + + +# Model info stubs keyed by model name +_MODEL_INFO = { + "gpt-3.5-turbo": { + "input_cost_per_token": 0.0000005, + "output_cost_per_token": 0.0000015, + "max_output_tokens": 4096, + "max_tokens": 4096, + "output_tokens_per_second": 50.0, + }, + "gpt-4o": { + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.0000100, + "max_output_tokens": 16384, + "max_tokens": 128000, + "output_tokens_per_second": 60.0, + }, + "claude-3-opus": { + "input_cost_per_token": 0.0000150, + "output_cost_per_token": 0.0000750, + "max_output_tokens": 4096, + "max_tokens": 200000, + "output_tokens_per_second": 20.0, + }, + "gpt-4o-mini": { + "input_cost_per_token": 0.00000015, + "output_cost_per_token": 0.0000006, + "max_output_tokens": 16384, + "max_tokens": 128000, + "output_tokens_per_second": 100.0, + }, +} + + +def _mock_get_model_info(model, **kwargs): + """Mock litellm.get_model_info using our test data.""" + if model in _MODEL_INFO: + return _MODEL_INFO[model] + raise Exception(f"Unknown model: {model}") + + +# --------------------------------------------------------------------------- +# _has_priorities +# --------------------------------------------------------------------------- + + +class TestHasPriorities: + def test_should_return_false_when_no_priorities_set(self): + prefs = _prefs() + assert _has_priorities(prefs) is False + + def test_should_return_false_when_all_zero(self): + prefs = _prefs(cost=0, speed=0, intelligence=0) + assert _has_priorities(prefs) is False + + def test_should_return_true_when_cost_set(self): + prefs = _prefs(cost=0.8) + assert _has_priorities(prefs) is True + + def test_should_return_true_when_intelligence_set(self): + prefs = _prefs(intelligence=0.5) + assert _has_priorities(prefs) is True + + +# --------------------------------------------------------------------------- +# _select_model_by_priority +# --------------------------------------------------------------------------- + + +class TestSelectModelByPriority: + """Tests for the priority-based scoring logic.""" + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + def test_should_prefer_cheapest_when_cost_priority_high(self, _mock): + """High costPriority should select the cheapest model.""" + prefs = _prefs(cost=1.0, speed=0, intelligence=0) + models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"] + result = _select_model_by_priority(models, prefs) + # gpt-4o-mini has the lowest combined cost + assert result == "gpt-4o-mini" + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + def test_should_prefer_smartest_when_intelligence_priority_high(self, _mock): + """High intelligencePriority should select the model with highest max_output_tokens.""" + prefs = _prefs(cost=0, speed=0, intelligence=1.0) + models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"] + result = _select_model_by_priority(models, prefs) + # gpt-4o and gpt-4o-mini both have 16384 max_output_tokens (tied) + # Either is acceptable + assert result in ("gpt-4o", "gpt-4o-mini") + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + def test_should_balance_cost_and_intelligence(self, _mock): + """Balanced priorities should pick a middle-ground model.""" + prefs = _prefs(cost=0.5, speed=0, intelligence=0.5) + models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"] + result = _select_model_by_priority(models, prefs) + # gpt-4o-mini is cheap AND has high max_output_tokens → best balance + assert result == "gpt-4o-mini" + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + def test_should_prefer_fastest_when_speed_priority_high(self, _mock): + """High speedPriority should prefer cheaper (faster proxy) models.""" + prefs = _prefs(cost=0, speed=1.0, intelligence=0) + models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"] + result = _select_model_by_priority(models, prefs) + # gpt-4o-mini has lowest cost → fastest proxy + assert result == "gpt-4o-mini" + + @patch( + "litellm.get_model_info", + side_effect=lambda m, **kw: (_ for _ in ()).throw(Exception("no info")), + ) + def test_should_return_none_when_no_model_info(self, _mock): + """If get_model_info fails for all models, return None.""" + prefs = _prefs(cost=1.0) + models = ["unknown-model-1", "unknown-model-2"] + result = _select_model_by_priority(models, prefs) + assert result is None + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + def test_should_handle_single_model(self, _mock): + """Single model should always be returned regardless of priorities.""" + prefs = _prefs(cost=1.0, intelligence=1.0) + result = _select_model_by_priority(["gpt-4o"], prefs) + assert result == "gpt-4o" + + def test_speed_priority_is_neutral_when_no_tps_data(self): + """When no candidate exposes output_tokens_per_second, speedPriority + must not fall back to context-window size as a latency proxy: that + biased selection toward the smallest-context model regardless of real + speed. With a neutral score the tie resolves to the first candidate, + so the larger-context model listed first is kept.""" + no_tps_info = { + "big-ctx": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "max_output_tokens": 100000, + "max_tokens": 100000, + }, + "small-ctx": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "max_output_tokens": 1000, + "max_tokens": 1000, + }, + } + + def info(model, **kwargs): + return no_tps_info[model] + + with patch("litellm.get_model_info", side_effect=info): + prefs = _prefs(speed=1.0) + # The inverse-max_output proxy would pick "small-ctx" here; a + # neutral score keeps the first candidate. + assert _select_model_by_priority(["big-ctx", "small-ctx"], prefs) == ( + "big-ctx" + ) + + +# --------------------------------------------------------------------------- +# _resolve_model_from_preferences — priority integration +# --------------------------------------------------------------------------- + + +class TestResolveModelPriorityIntegration: + """End-to-end tests for priority selection within _resolve_model_from_preferences.""" + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch( + "litellm.model_list", + [ + {"model_name": "gpt-3.5-turbo"}, + {"model_name": "gpt-4o"}, + {"model_name": "gpt-4o-mini"}, + ], + ) + def test_should_use_priority_when_hints_empty(self, _mock_info): + """With no hints but priorities set, should use priority-based selection.""" + prefs = _prefs(cost=1.0, speed=0, intelligence=0) + result = _resolve_model_from_preferences(prefs, default_model="gpt-4o") + # Should pick cheapest, NOT fall through to default_model + assert result == "gpt-4o-mini" + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch( + "litellm.model_list", + [ + {"model_name": "gpt-3.5-turbo"}, + {"model_name": "gpt-4o"}, + {"model_name": "gpt-4o-mini"}, + ], + ) + def test_should_skip_priority_when_no_priorities_set(self, _mock_info): + """With no priorities set, should fall through to default_model.""" + prefs = _prefs() # no priorities + result = _resolve_model_from_preferences(prefs, default_model="gpt-4o") + assert result == "gpt-4o" + + @patch("litellm.get_model_info", side_effect=_mock_get_model_info) + @patch("litellm.proxy.proxy_server.llm_router", None) + @patch( + "litellm.model_list", + [ + {"model_name": "gpt-3.5-turbo"}, + {"model_name": "gpt-4o"}, + {"model_name": "gpt-4o-mini"}, + ], + ) + def test_should_prefer_hint_over_priority(self, _mock_info): + """Hints should take precedence over priority-based selection.""" + hints = [SimpleNamespace(name="gpt-4o")] + prefs = _prefs(hints=hints, cost=1.0) # cost says cheap, but hint says gpt-4o + result = _resolve_model_from_preferences(prefs, default_model="gpt-3.5-turbo") + assert result == "gpt-4o" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py new file mode 100644 index 00000000000..d5c636baead --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py @@ -0,0 +1,147 @@ +""" +Tests for _build_sampling_request header forwarding. + +Verifies that the synthetic FastAPI Request built for sampling sub-calls +correctly propagates the original MCP connection's headers and client IP +so that header-dependent guardrails, routing hooks, and trace correlation +function correctly. +""" + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _build_sampling_request, +) + + +class TestBuildSamplingRequest: + """Tests for the _build_sampling_request helper.""" + + def test_should_include_content_type_by_default(self): + """Even with no raw headers, content-type must be present.""" + req = _build_sampling_request() + headers = dict(req.headers) + assert headers.get("content-type") == "application/json" + + def test_should_forward_raw_headers(self): + """Headers from the original MCP connection should be forwarded.""" + raw = { + "x-litellm-tags": "tag1,tag2", + "x-litellm-trace-id": "trace-abc-123", + "user-agent": "MCP-Client/1.0", + "authorization": "Bearer sk-test", + } + req = _build_sampling_request(raw_headers=raw) + headers = dict(req.headers) + + assert headers.get("x-litellm-tags") == "tag1,tag2" + assert headers.get("x-litellm-trace-id") == "trace-abc-123" + assert headers.get("user-agent") == "MCP-Client/1.0" + assert headers.get("authorization") == "Bearer sk-test" + + def test_should_skip_hop_by_hop_headers(self): + """content-length and transfer-encoding should not be forwarded.""" + raw = { + "content-length": "42", + "transfer-encoding": "chunked", + "x-custom": "keep-me", + } + req = _build_sampling_request(raw_headers=raw) + headers = dict(req.headers) + + assert "content-length" not in headers + assert "transfer-encoding" not in headers + assert headers.get("x-custom") == "keep-me" + + def test_should_not_duplicate_content_type(self): + """If raw_headers includes content-type, don't add it twice.""" + raw = {"content-type": "text/plain"} + req = _build_sampling_request(raw_headers=raw) + # Count how many content-type headers are present + ct_count = sum(1 for k, _ in req.scope["headers"] if k == b"content-type") + assert ct_count == 1 + + def test_should_inject_client_ip_as_x_forwarded_for(self): + """client_ip should be injected as x-forwarded-for.""" + req = _build_sampling_request(client_ip="10.0.0.42") + headers = dict(req.headers) + assert headers.get("x-forwarded-for") == "10.0.0.42" + + def test_should_not_override_existing_x_forwarded_for(self): + """Caller-supplied x-forwarded-for is stripped; resolved client_ip wins.""" + raw = {"x-forwarded-for": "192.168.1.1"} + req = _build_sampling_request(raw_headers=raw, client_ip="10.0.0.42") + headers = dict(req.headers) + assert headers.get("x-forwarded-for") == "10.0.0.42" + + def test_should_set_correct_path(self): + """The synthetic request should have the sampling path.""" + req = _build_sampling_request() + assert req.scope["path"] == "/mcp/sampling/createMessage" + + def test_server_should_default_to_litellm_port(self): + """Server tuple should use port 4000 (LiteLLM default), not 0.""" + req = _build_sampling_request() + _host, _port = req.scope["server"] + assert _port == 4000, f"Expected default LiteLLM port 4000, got {_port}" + + def test_should_populate_client_tuple_from_client_ip(self): + """request.client.host must return the real client IP for + IP-based routing and guardrails.""" + req = _build_sampling_request(client_ip="10.0.0.42") + assert req.scope.get("client") is not None + assert req.scope["client"][0] == "10.0.0.42" + # Verify request.client.host works (Starlette Address) + assert req.client is not None + assert req.client.host == "10.0.0.42" + + def test_should_not_set_client_when_no_ip(self): + """If no client_ip is provided, client should not be in scope.""" + req = _build_sampling_request() + assert "client" not in req.scope + + def test_should_skip_all_hop_by_hop_headers(self): + """All hop-by-hop headers must be filtered, not just content-length + and transfer-encoding.""" + raw = { + "content-length": "42", + "transfer-encoding": "chunked", + "connection": "keep-alive", + "keep-alive": "timeout=5", + "upgrade": "websocket", + "te": "trailers", + "trailer": "Expires", + "x-custom": "keep-me", + } + req = _build_sampling_request(raw_headers=raw) + headers = dict(req.headers) + + for hop_header in [ + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "te", + "trailer", + ]: + assert ( + hop_header not in headers + ), f"Hop-by-hop header '{hop_header}' should be filtered" + assert headers.get("x-custom") == "keep-me" + + def test_should_forward_traceparent_header(self): + """traceparent header must be forwarded for trace correlation.""" + raw = { + "traceparent": "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01", + } + req = _build_sampling_request(raw_headers=raw) + headers = dict(req.headers) + assert headers.get("traceparent") == ( + "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01" + ) + + def test_should_forward_x_litellm_api_key(self): + """x-litellm-api-key header must be forwarded for auth.""" + raw = {"x-litellm-api-key": "sk-proxy-key-123"} + req = _build_sampling_request(raw_headers=raw) + headers = dict(req.headers) + assert headers.get("x-litellm-api-key") == "sk-proxy-key-123" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py new file mode 100644 index 00000000000..bb17a8f7104 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py @@ -0,0 +1,180 @@ +""" +Tests for MCP sampling handler response/tool conversion. + +Covers the translation of a LiteLLM completion response back into MCP +`CreateMessageResult` / `CreateMessageResultWithTools`, plus the helpers that +convert MCP tool definitions, tool-choice modes, and image/audio content into +OpenAI request format. +""" + +import json +from types import SimpleNamespace + +from mcp.types import ( + CreateMessageResult, + CreateMessageResultWithTools, + ErrorData, + TextContent, + ToolUseContent, +) + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_content_to_openai, + _convert_mcp_tool_choice_to_openai, + _convert_mcp_tools_to_openai, + _convert_openai_response_to_mcp_result, + _convert_single_content, +) + + +def _tool_call(*, call_id: str, name: str, arguments): + return SimpleNamespace( + id=call_id, function=SimpleNamespace(name=name, arguments=arguments) + ) + + +def _response(*, content=None, tool_calls=None, finish_reason="stop", model="gpt-4o"): + message = SimpleNamespace(content=content, tool_calls=tool_calls) + choice = SimpleNamespace(message=message, finish_reason=finish_reason) + return SimpleNamespace(choices=[choice], model=model) + + +class TestConvertOpenAIResponseToMcpResult: + def test_should_return_error_data_when_no_choices(self): + response = SimpleNamespace(choices=[], model="gpt-4o") + result = _convert_openai_response_to_mcp_result(response, "gpt-4o") + assert isinstance(result, ErrorData) + assert "no choices" in result.message.lower() + + def test_should_convert_plain_text_response(self): + result = _convert_openai_response_to_mcp_result( + _response(content="hello world"), "gpt-4o" + ) + assert isinstance(result, CreateMessageResult) + assert isinstance(result.content, TextContent) + assert result.content.text == "hello world" + assert result.role == "assistant" + assert result.stopReason == "endTurn" + + def test_should_map_length_finish_reason_to_max_tokens(self): + result = _convert_openai_response_to_mcp_result( + _response(content="truncated", finish_reason="length"), "gpt-4o" + ) + assert result.stopReason == "maxTokens" + + def test_should_prefer_actual_model_from_response(self): + result = _convert_openai_response_to_mcp_result( + _response(content="hi", model="gpt-4o-2024-08-06"), "gpt-4o" + ) + assert result.model == "gpt-4o-2024-08-06" + + def test_should_convert_tool_calls_response(self): + tc = _tool_call( + call_id="call_1", + name="get_weather", + arguments=json.dumps({"city": "NYC"}), + ) + result = _convert_openai_response_to_mcp_result( + _response(content=None, tool_calls=[tc], finish_reason="tool_calls"), + "gpt-4o", + ) + assert isinstance(result, CreateMessageResultWithTools) + assert result.stopReason == "toolUse" + tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)] + assert len(tool_uses) == 1 + assert tool_uses[0].name == "get_weather" + assert tool_uses[0].id == "call_1" + assert tool_uses[0].input == {"city": "NYC"} + + def test_should_keep_text_alongside_tool_calls(self): + tc = _tool_call(call_id="call_1", name="search", arguments="{}") + result = _convert_openai_response_to_mcp_result( + _response( + content="let me check", tool_calls=[tc], finish_reason="tool_calls" + ), + "gpt-4o", + ) + texts = [c for c in result.content if isinstance(c, TextContent)] + assert texts and texts[0].text == "let me check" + + def test_should_wrap_unparsable_tool_arguments_as_raw(self): + tc = _tool_call(call_id="call_1", name="bad", arguments="not-json{") + result = _convert_openai_response_to_mcp_result( + _response(tool_calls=[tc], finish_reason="tool_calls"), "gpt-4o" + ) + tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)] + assert tool_uses[0].input == {"raw": "not-json{"} + + +class TestConvertMcpToolsToOpenAI: + def test_should_return_none_when_no_tools(self): + assert _convert_mcp_tools_to_openai(None) is None + + def test_should_convert_tool_with_schema(self): + schema = {"type": "object", "properties": {"q": {"type": "string"}}} + tool = SimpleNamespace( + name="search", description="search the web", inputSchema=schema + ) + result = _convert_mcp_tools_to_openai([tool]) + assert result == [ + { + "type": "function", + "function": { + "name": "search", + "description": "search the web", + "parameters": schema, + }, + } + ] + + def test_should_default_description_and_parameters(self): + tool = SimpleNamespace(name="noop", description=None, inputSchema=None) + result = _convert_mcp_tools_to_openai([tool]) + fn = result[0]["function"] + assert fn["description"] == "" + assert fn["parameters"] == {"type": "object", "properties": {}} + + +class TestConvertMcpToolChoiceToOpenAI: + def test_should_return_none_when_no_choice(self): + assert _convert_mcp_tool_choice_to_openai(None) is None + + def test_should_map_known_modes(self): + for mode in ("auto", "required", "none"): + choice = SimpleNamespace(mode=mode) + assert _convert_mcp_tool_choice_to_openai(choice) == mode + + def test_should_default_unknown_mode_to_auto(self): + choice = SimpleNamespace(mode="banana") + assert _convert_mcp_tool_choice_to_openai(choice) == "auto" + + +class TestConvertImageAndAudioContent: + def test_should_convert_image_to_data_uri(self): + content = SimpleNamespace(type="image", data="aGVsbG8=", mimeType="image/jpeg") + result = _convert_single_content(content) + assert result == { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,aGVsbG8="}, + } + + def test_should_map_audio_mime_to_format(self): + content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/mp3") + result = _convert_single_content(content) + assert result["type"] == "input_audio" + assert result["input_audio"] == {"data": "Zm9v", "format": "mp3"} + + def test_should_default_unknown_audio_mime_to_wav(self): + content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/weird") + result = _convert_single_content(content) + assert result["input_audio"]["format"] == "wav" + + def test_should_flatten_list_content(self): + items = [ + SimpleNamespace(type="text", text="a"), + SimpleNamespace(type="image", data="x", mimeType="image/png"), + ] + result = _convert_mcp_content_to_openai(items) + assert isinstance(result, list) + assert result[0] == {"type": "text", "text": "a"} + assert result[1]["type"] == "image_url" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py new file mode 100644 index 00000000000..b4b219e958c --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py @@ -0,0 +1,312 @@ +""" +Tests for MCP sampling handler tool_use / tool_result content conversion. + +Verifies that multi-turn tool-calling conversations from upstream MCP +servers are faithfully converted to OpenAI format instead of being +reduced to lossy plain-text stubs. +""" + +import json +from types import SimpleNamespace +from typing import Any, Dict + +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_messages_to_openai, + _convert_single_content, +) + + +# --------------------------------------------------------------------------- +# Helpers — lightweight MCP type stand-ins +# --------------------------------------------------------------------------- + + +def _text(text: str) -> SimpleNamespace: + return SimpleNamespace(type="text", text=text) + + +def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleNamespace: + return SimpleNamespace(type="tool_use", name=name, id=tool_id, input=input_data) + + +def _tool_result( + *, tool_use_id: str, content: Any = None, is_error: bool = False +) -> SimpleNamespace: + if content is None: + content = [] + return SimpleNamespace( + type="tool_result", toolUseId=tool_use_id, content=content, isError=is_error + ) + + +def _sampling_msg(role: str, content: Any) -> SimpleNamespace: + return SimpleNamespace(role=role, content=content) + + +# --------------------------------------------------------------------------- +# _convert_single_content — tool_use +# --------------------------------------------------------------------------- + + +class TestConvertSingleContentToolUse: + """Tests for the tool_use branch of _convert_single_content.""" + + def test_should_produce_function_call_dict(self): + """tool_use must produce a proper function-call dict, not a text stub.""" + tu = _tool_use(name="get_weather", tool_id="call_123", input_data={"city": "NYC"}) + result = _convert_single_content(tu) + + assert result["_marker_type"] == "tool_use" + assert result["type"] == "function" + assert result["id"] == "call_123" + assert result["function"]["name"] == "get_weather" + assert json.loads(result["function"]["arguments"]) == {"city": "NYC"} + + def test_should_not_produce_text_stub(self): + """Regression: the old code produced '[Tool call: get_weather]'.""" + tu = _tool_use(name="get_weather", tool_id="call_1", input_data={}) + result = _convert_single_content(tu) + + # Must NOT be a text content part + assert result.get("type") != "text" + assert "Tool call" not in str(result) + + def test_should_handle_empty_input(self): + tu = _tool_use(name="no_args_tool", tool_id="call_2", input_data={}) + result = _convert_single_content(tu) + + assert json.loads(result["function"]["arguments"]) == {} + + +# --------------------------------------------------------------------------- +# _convert_single_content — tool_result +# --------------------------------------------------------------------------- + + +class TestConvertSingleContentToolResult: + """Tests for the tool_result branch of _convert_single_content.""" + + def test_should_produce_tool_role_message(self): + """tool_result must produce a tool-role dict, not a text content part.""" + tr = _tool_result( + tool_use_id="call_123", + content=[_text("Temperature: 72°F")], + ) + result = _convert_single_content(tr) + + assert result["_marker_type"] == "tool_result" + assert result["role"] == "tool" + assert result["tool_call_id"] == "call_123" + assert "72°F" in result["content"] + + def test_should_handle_empty_content(self): + tr = _tool_result(tool_use_id="call_456", content=[]) + result = _convert_single_content(tr) + + assert result["role"] == "tool" + assert result["tool_call_id"] == "call_456" + assert result["content"] == "" + + def test_should_concatenate_multiple_text_parts(self): + tr = _tool_result( + tool_use_id="call_789", + content=[_text("Line 1"), _text("Line 2")], + ) + result = _convert_single_content(tr) + assert "Line 1" in result["content"] + assert "Line 2" in result["content"] + + +# --------------------------------------------------------------------------- +# _convert_mcp_messages_to_openai — multi-turn tool calling +# --------------------------------------------------------------------------- + + +class TestConvertMcpMessagesMultiTurnTools: + """End-to-end tests for multi-turn tool-calling message sequences.""" + + def test_should_convert_assistant_tool_use_to_tool_calls_array(self): + """An assistant message with tool_use content should produce + a proper tool_calls array, not a text stub.""" + messages = [ + _sampling_msg("assistant", _tool_use( + name="search", tool_id="call_1", input_data={"query": "LiteLLM"} + )), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "assistant" + assert "tool_calls" in msg + assert len(msg["tool_calls"]) == 1 + tc = msg["tool_calls"][0] + assert tc["function"]["name"] == "search" + assert tc["id"] == "call_1" + + def test_should_convert_user_tool_result_to_tool_role_message(self): + """A user message with tool_result content should produce + a separate role='tool' message.""" + messages = [ + _sampling_msg("user", _tool_result( + tool_use_id="call_1", + content=[_text("Found 42 results")], + )), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "tool" + assert msg["tool_call_id"] == "call_1" + assert "42 results" in msg["content"] + + def test_should_handle_full_tool_calling_round_trip(self): + """Simulate a complete tool-calling conversation: + user → assistant(tool_use) → user(tool_result) → assistant(text) + """ + messages = [ + _sampling_msg("user", _text("What's the weather in NYC?")), + _sampling_msg("assistant", _tool_use( + name="get_weather", tool_id="call_w1", + input_data={"city": "NYC"}, + )), + _sampling_msg("user", _tool_result( + tool_use_id="call_w1", + content=[_text("72°F, sunny")], + )), + _sampling_msg("assistant", _text("It's 72°F and sunny in NYC!")), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 4 + + # 1. User message + assert result[0]["role"] == "user" + + # 2. Assistant with tool_calls + assert result[1]["role"] == "assistant" + assert "tool_calls" in result[1] + assert result[1]["tool_calls"][0]["function"]["name"] == "get_weather" + + # 3. Tool result + assert result[2]["role"] == "tool" + assert result[2]["tool_call_id"] == "call_w1" + + # 4. Final assistant text + assert result[3]["role"] == "assistant" + assert "72°F" in str(result[3]["content"]) + + def test_should_handle_mixed_text_and_tool_use_in_assistant(self): + """An assistant message with both text and tool_use content.""" + messages = [ + _sampling_msg("assistant", [ + _text("Let me check that for you."), + _tool_use(name="lookup", tool_id="call_lu1", input_data={"id": 42}), + ]), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "assistant" + assert "tool_calls" in msg + assert msg["tool_calls"][0]["function"]["name"] == "lookup" + # Text content should also be present + assert msg.get("content") is not None + + def test_should_handle_multiple_tool_uses_in_single_message(self): + """Multiple tool_use items in a single assistant message → multiple tool_calls.""" + messages = [ + _sampling_msg("assistant", [ + _tool_use(name="tool_a", tool_id="call_a", input_data={}), + _tool_use(name="tool_b", tool_id="call_b", input_data={"x": 1}), + ]), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + msg = result[0] + assert len(msg["tool_calls"]) == 2 + names = {tc["function"]["name"] for tc in msg["tool_calls"]} + assert names == {"tool_a", "tool_b"} + + def test_should_handle_multiple_tool_results_in_single_message(self): + """Multiple tool_result items in a single user message → multiple tool messages.""" + messages = [ + _sampling_msg("user", [ + _tool_result(tool_use_id="call_a", content=[_text("Result A")]), + _tool_result(tool_use_id="call_b", content=[_text("Result B")]), + ]), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 2 + assert all(m["role"] == "tool" for m in result) + ids = {m["tool_call_id"] for m in result} + assert ids == {"call_a", "call_b"} + + def test_should_preserve_system_prompt(self): + """System prompt should still be emitted first.""" + messages = [_sampling_msg("user", _text("Hi"))] + result = _convert_mcp_messages_to_openai( + messages, system_prompt="You are helpful." + ) + + assert result[0]["role"] == "system" + assert result[0]["content"] == "You are helpful." + + +# --------------------------------------------------------------------------- +# _convert_mcp_messages_to_openai — marker hoisting on unexpected roles +# --------------------------------------------------------------------------- + + +class TestConvertMcpMessagesMarkerHoisting: + """The role-matched fast paths only fire for assistant/tool_use and + user/tool_result. Content that arrives on an unexpected role must still + be hoisted to the correct message position by the generic fallback, + not silently dropped or embedded inline as a content part.""" + + def test_should_hoist_tool_use_arriving_on_user_role(self): + messages = [ + _sampling_msg("user", _tool_use( + name="search", tool_id="call_1", input_data={"q": "x"} + )), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + assert result[0]["role"] == "assistant" + assert result[0]["tool_calls"][0]["function"]["name"] == "search" + + def test_should_hoist_tool_result_arriving_on_assistant_role(self): + messages = [ + _sampling_msg("assistant", _tool_result( + tool_use_id="call_1", content=[_text("done")] + )), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + assert result[0]["role"] == "tool" + assert result[0]["tool_call_id"] == "call_1" + assert "done" in result[0]["content"] + + def test_should_keep_text_when_hoisting_tool_use_on_user_role(self): + messages = [ + _sampling_msg("user", [ + _text("here you go"), + _tool_use(name="lookup", tool_id="call_2", input_data={}), + ]), + ] + result = _convert_mcp_messages_to_openai(messages) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "assistant" + assert msg["tool_calls"][0]["function"]["name"] == "lookup" + assert any( + isinstance(p, dict) and p.get("text") == "here you go" + for p in msg["content"] + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 6b6c7bc37d5..227cf3f4bcf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,5 +1,4 @@ import asyncio -import contextlib import contextvars from datetime import datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch @@ -894,7 +893,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): extra_headers=None, add_prefix=True, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): if server.name == "working_server": # Working server returns tools @@ -1000,7 +999,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): extra_headers=None, add_prefix=True, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -1122,8 +1121,8 @@ async def test_concurrent_initialize_session_managers(): # Reset state before test original_initialized = mcp_server._SESSION_MANAGERS_INITIALIZED original_session_cm = mcp_server._session_manager_cm - original_session_stateful_cm = mcp_server._session_manager_stateful_cm - original_sse_session_cm = mcp_server._sse_session_manager_cm + original_stateful_cm = mcp_server._session_manager_stateful_cm + original_sse_cm = mcp_server._sse_session_manager_cm original_cleanup_task = mcp_server._stateful_auth_context_cleanup_task try: @@ -1131,30 +1130,38 @@ async def test_concurrent_initialize_session_managers(): mcp_server._session_manager_cm = None mcp_server._session_manager_stateful_cm = None mcp_server._sse_session_manager_cm = None - mcp_server._stateful_auth_context_cleanup_task = None - # Mock the session managers to avoid actual MCP initialization + # Create mock context managers for all three session managers + mock_cm_stateless = AsyncMock() + mock_cm_stateless.__aenter__ = AsyncMock() + mock_cm_stateless.__aexit__ = AsyncMock() + + mock_cm_stateful = AsyncMock() + mock_cm_stateful.__aenter__ = AsyncMock() + mock_cm_stateful.__aexit__ = AsyncMock() + + mock_cm_sse = AsyncMock() + mock_cm_sse.__aenter__ = AsyncMock() + mock_cm_sse.__aexit__ = AsyncMock() + with ( - patch( - "litellm.proxy._experimental.mcp_server.server.session_manager_stateless" - ) as mock_session_manager_stateless, - patch( - "litellm.proxy._experimental.mcp_server.server.session_manager_stateful" - ) as mock_session_manager_stateful, - patch( - "litellm.proxy._experimental.mcp_server.server.sse_session_manager" - ) as mock_sse_session_manager, + patch.object( + mcp_server.session_manager_stateless, + "run", + return_value=mock_cm_stateless, + ) as mock_stateless_run, + patch.object( + mcp_server.session_manager_stateful, + "run", + return_value=mock_cm_stateful, + ) as mock_stateful_run, + patch.object( + mcp_server.sse_session_manager, + "run", + return_value=mock_cm_sse, + ) as mock_sse_run, patch("litellm.proxy._experimental.mcp_server.server.verbose_logger"), ): - # Mock the run() method to return a mock context manager - mock_cm = AsyncMock() - mock_cm.__aenter__ = AsyncMock() - mock_cm.__aexit__ = AsyncMock() - - mock_session_manager_stateless.run.return_value = mock_cm - mock_session_manager_stateful.run.return_value = mock_cm - mock_sse_session_manager.run.return_value = mock_cm - # Create multiple concurrent tasks that call initialize_session_managers async def init_task(): await initialize_session_managers() @@ -1171,19 +1178,25 @@ async def test_concurrent_initialize_session_managers(): # Each session manager.run() should only be called once due to the lock assert ( - mock_session_manager_stateless.run.call_count == 1 - ), f"Expected 1 call to session_manager_stateless.run(), got {mock_session_manager_stateless.run.call_count}" + mock_stateless_run.call_count == 1 + ), f"Expected 1 call to session_manager_stateless.run(), got {mock_stateless_run.call_count}" assert ( - mock_session_manager_stateful.run.call_count == 1 - ), f"Expected 1 call to session_manager_stateful.run(), got {mock_session_manager_stateful.run.call_count}" + mock_stateful_run.call_count == 1 + ), f"Expected 1 call to session_manager_stateful.run(), got {mock_stateful_run.call_count}" assert ( - mock_sse_session_manager.run.call_count == 1 - ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_session_manager.run.call_count}" + mock_sse_run.call_count == 1 + ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_run.call_count}" - # The context managers should only be entered once each (3 managers) + # The context managers should only be entered once each assert ( - mock_cm.__aenter__.call_count == 3 - ), f"Expected 3 calls to __aenter__ (one per session manager), got {mock_cm.__aenter__.call_count}" + mock_cm_stateless.__aenter__.call_count == 1 + ), f"Expected 1 call to stateless __aenter__, got {mock_cm_stateless.__aenter__.call_count}" + assert ( + mock_cm_stateful.__aenter__.call_count == 1 + ), f"Expected 1 call to stateful __aenter__, got {mock_cm_stateful.__aenter__.call_count}" + assert ( + mock_cm_sse.__aenter__.call_count == 1 + ), f"Expected 1 call to sse __aenter__, got {mock_cm_sse.__aenter__.call_count}" # State should be properly set assert mcp_server._SESSION_MANAGERS_INITIALIZED is True @@ -1195,14 +1208,12 @@ async def test_concurrent_initialize_session_managers(): leaked_task = mcp_server._stateful_auth_context_cleanup_task if leaked_task is not None and leaked_task is not original_cleanup_task: leaked_task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await leaked_task # Restore original state mcp_server._SESSION_MANAGERS_INITIALIZED = original_initialized mcp_server._session_manager_cm = original_session_cm - mcp_server._session_manager_stateful_cm = original_session_stateful_cm - mcp_server._sse_session_manager_cm = original_sse_session_cm + mcp_server._session_manager_stateful_cm = original_stateful_cm + mcp_server._sse_session_manager_cm = original_sse_cm mcp_server._stateful_auth_context_cleanup_task = original_cleanup_task @@ -1637,10 +1648,7 @@ async def test_mcp_routing_initialize_rejected_when_owner_at_session_cap(): active = {f"s{i}": 1 for i in range(cap)} # all in flight -> cannot evict contexts = {f"s{i}": MagicMock() for i in range(cap)} - init_body = ( - b'{"jsonrpc":"2.0","id":1,"method":"initialize",' - b'"params":{"protocolVersion":"2024-11-05"}}' - ) + init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}' scope = { "type": "http", "method": "POST", @@ -2587,6 +2595,134 @@ async def test_stateful_mcp_get_stream_does_not_block_post(): mcp_server._stateful_session_locks.pop(session_id, None) +def test_jsonrpc_text_has_top_level_method_ignores_nested_method(): + """The top-level-key scan must not be fooled by a ``method`` field nested + inside a JSON-RPC response's ``result`` payload — a flat substring search + would, and that misread is what deadlocks the session lock.""" + from litellm.proxy._experimental.mcp_server.server import ( + _jsonrpc_text_has_top_level_method, + ) + + request = '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}' + assert _jsonrpc_text_has_top_level_method(request) is True + + # method key out of order (after params) is still top-level + reordered = '{"jsonrpc":"2.0","params":{"x":1},"method":"foo"}' + assert _jsonrpc_text_has_top_level_method(reordered) is True + + # response whose result nests a "method" key (and arrays of them) + response = ( + '{"jsonrpc":"2.0","id":1,"result":{"toolResult":{"method":"GET"},' + '"steps":[{"method":"x"}]}}' + ) + assert _jsonrpc_text_has_top_level_method(response) is False + + # truncated response: result value never closes, no top-level method seen + truncated = '{"jsonrpc":"2.0","id":1,"result":{"text":"' + "q" * 5000 + assert _jsonrpc_text_has_top_level_method(truncated) is False + + +@pytest.mark.asyncio +async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): + """Regression: a large JSON-RPC *response* POST whose ``result`` payload + nests a ``method`` key must skip the per-session lock so it does not + deadlock behind the in-flight request POST that is holding the lock while + it awaits this very response (e.g. sampling/createMessage).""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "nested-method-response-session" + owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner") + mcp_server._stateful_session_auth_contexts[session_id] = ( + mcp_server.MCPAuthenticatedUser(user_api_key_auth=owner_auth) + ) + mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for( + owner_auth + ) + + gate = asyncio.Event() + request_in_handle = asyncio.Event() + response_handled = asyncio.Event() + + async def handle(s, r, se): + msg = await r() + body = msg.get("body", b"") or b"" + if b'"result"' in body: + response_handled.set() + else: + request_in_handle.set() + await gate.wait() + + async def call(body: bytes): + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"mcp-session-id", session_id.encode())], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": body, + "more_body": False, + } + ) + await handle_streamable_http_mcp(scope, receive, AsyncMock()) + + # The in-flight request POST holds the session lock while blocked. + request_body = b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}' + # A JSON-RPC response larger than the routing peek cap so it can't be fully + # parsed, with a nested "method" key in the first bytes to trip a flat + # substring heuristic. + response_body = ( + '{"jsonrpc":"2.0","id":99,"result":{"toolResult":' + '{"method":"GET","payload":"' + ("x" * 5000) + '"}}}' + ).encode() + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, "handle_request", side_effect=handle + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + req_task = asyncio.create_task(call(request_body)) + await asyncio.wait_for(request_in_handle.wait(), timeout=1.0) + + resp_task = asyncio.create_task(call(response_body)) + # Under a flat substring heuristic the response would acquire the + # lock held by req_task and this wait would time out (deadlock). + await asyncio.wait_for(response_handled.wait(), timeout=1.0) + + gate.set() + await asyncio.gather(req_task, resp_task) + finally: + gate.set() + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_owners.pop(session_id, None) + mcp_server._stateful_session_locks.pop(session_id, None) + mcp_server._stateful_session_active_request_counts.pop(session_id, None) + + @pytest.mark.asyncio @pytest.mark.no_parallel async def test_mcp_routing_with_conflicting_alias_and_group_name(): @@ -2729,7 +2865,7 @@ async def test_oauth2_headers_passed_to_mcp_client(): mcp_auth_header=None, extra_headers=None, stdio_env=None, - subject_token=None, + **kwargs, ): # Capture the arguments for verification captured_client_args.update( @@ -2738,7 +2874,7 @@ async def test_oauth2_headers_passed_to_mcp_client(): "mcp_auth_header": mcp_auth_header, "extra_headers": extra_headers, "stdio_env": stdio_env, - "subject_token": subject_token, + "kwargs": kwargs, } ) # Return a mock client that doesn't actually connect @@ -2764,6 +2900,16 @@ async def test_oauth2_headers_passed_to_mcp_client(): "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", AsyncMock(return_value=[oauth2_server]), ), + patch( + "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user", + new_callable=AsyncMock, + return_value={}, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new_callable=AsyncMock, + return_value=None, + ), ): # Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client await _get_tools_from_mcp_servers( @@ -2840,7 +2986,7 @@ async def test_list_tools_single_server_unprefixed_names(): extra_headers=None, add_prefix=False, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" @@ -2922,7 +3068,7 @@ async def test_list_tools_multiple_servers_prefixed_names(): extra_headers=None, add_prefix=True, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): tool = MagicMock() # When multiple servers, add_prefix should be True -> prefixed names @@ -3189,7 +3335,7 @@ async def test_list_tools_filters_by_key_team_permissions(): extra_headers=None, add_prefix=False, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): # Return 4 tools, but only 2 should be allowed tool1 = MagicMock() @@ -3299,7 +3445,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): extra_headers=None, add_prefix=False, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): # Return 4 tools tool1 = MagicMock() @@ -3395,7 +3541,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): extra_headers=None, add_prefix=False, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): # Return 3 tools tool1 = MagicMock() @@ -3494,7 +3640,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): extra_headers=None, add_prefix=True, raw_headers=None, - user_api_key_auth=None, + **kwargs, ): # Return tools WITH prefix (as they come from MCP server) tool1 = MagicMock() @@ -5178,3 +5324,42 @@ def test_get_forwarded_auth_from_scope_skips_when_no_litellm_key_header(): ] } assert _get_forwarded_auth_from_scope(scope) is None + + +@pytest.mark.asyncio +async def test_create_mcp_client_sampling_disabled_by_default(): + """Sampling callback must be None when allow_sampling is not set (default False).""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + server = MCPServer( + server_id="no-sampling", + name="no-sampling", + url="https://example.com/mcp", + transport=MCPTransport.http, + ) + + client = await manager._create_mcp_client(server=server) + assert client._sampling_callback is None + + +@pytest.mark.asyncio +async def test_create_mcp_client_sampling_enabled(): + """Sampling callback must be set when allow_sampling=True.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + server = MCPServer( + server_id="with-sampling", + name="with-sampling", + url="https://example.com/mcp", + transport=MCPTransport.http, + allow_sampling=True, + ) + + client = await manager._create_mcp_client(server=server) + assert client._sampling_callback is not None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 32e4ec19311..e7d0ee6247b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -320,9 +320,7 @@ class TestMCPServerManager: async def mock_get_tools_from_server( server, mcp_auth_header=None, - mcp_protocol_version=None, - raw_headers=None, - user_api_key_auth=None, + **kwargs, ): if server.name == "github": tool1 = MagicMock() @@ -375,9 +373,7 @@ class TestMCPServerManager: async def mock_get_tools_from_server( server, mcp_auth_header=None, - mcp_protocol_version=None, - raw_headers=None, - user_api_key_auth=None, + **kwargs, ): assert mcp_auth_header == "legacy-token" # Should use legacy header tool = MagicMock() @@ -414,9 +410,7 @@ class TestMCPServerManager: async def mock_get_tools_from_server( server, mcp_auth_header=None, - mcp_protocol_version=None, - raw_headers=None, - user_api_key_auth=None, + **kwargs, ): assert ( mcp_auth_header == "server-specific-token" @@ -457,7 +451,7 @@ class TestMCPServerManager: captured_extra_headers = None async def capture_create_mcp_client( - server, mcp_auth_header, extra_headers, stdio_env, subject_token=None + server, mcp_auth_header, extra_headers, stdio_env, **kwargs ): # pragma: no cover - helper nonlocal captured_extra_headers captured_extra_headers = extra_headers @@ -507,7 +501,7 @@ class TestMCPServerManager: captured_extra_headers = None async def capture_create_mcp_client( - server, mcp_auth_header, extra_headers, stdio_env, subject_token=None + server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs ): # pragma: no cover - helper nonlocal captured_extra_headers captured_extra_headers = extra_headers @@ -560,7 +554,7 @@ class TestMCPServerManager: captured_extra_headers = None async def capture_create_mcp_client( - server, mcp_auth_header, extra_headers, stdio_env, subject_token=None + server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs ): # pragma: no cover - helper nonlocal captured_extra_headers captured_extra_headers = extra_headers @@ -616,7 +610,7 @@ class TestMCPServerManager: captured_extra_headers = None async def capture_create_mcp_client( - server, mcp_auth_header, extra_headers, stdio_env, subject_token=None + server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs ): # pragma: no cover - helper nonlocal captured_extra_headers captured_extra_headers = extra_headers @@ -1166,9 +1160,7 @@ class TestMCPServerManager: async def mock_get_tools_from_server( server, mcp_auth_header=None, - mcp_protocol_version=None, - raw_headers=None, - user_api_key_auth=None, + **kwargs, ): assert ( mcp_auth_header == "server-specific-token" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 549afd774b0..d52af94c47f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -9,8 +9,6 @@ they may send a stale `mcp-session-id` header. This test verifies that: import asyncio from unittest.mock import AsyncMock, MagicMock, patch - -from fastapi import HTTPException from litellm.types.mcp import MCPAuth import pytest @@ -600,6 +598,8 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): Per-user OAuth server with no stored token should fail fast with 401 + WWW-Authenticate so PKCE can start. """ + from fastapi import HTTPException + try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, @@ -612,8 +612,13 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): "type": "http", "method": "POST", "path": "/mcp", + "scheme": "http", + "query_string": b"", + "root_path": "", + "server": ("localhost", 8000), "headers": [ (b"content-type", b"application/json"), + (b"host", b"localhost:8000"), ], } receive = AsyncMock() @@ -660,11 +665,12 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): with pytest.raises(HTTPException) as exc_info: await handle_streamable_http_mcp(scope, receive, send) - exc = exc_info.value - assert exc.status_code == 401 - assert "www-authenticate" in exc.headers + # Verify a 401 was raised assert mock_get_stored_token.await_count == 1 assert mock_handle_request.await_count == 0 + assert exc_info.value.status_code == 401 + assert "www-authenticate" in exc_info.value.headers + assert "Bearer authorization_uri=" in exc_info.value.headers["www-authenticate"] @pytest.mark.asyncio @@ -685,11 +691,22 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): "type": "http", "method": "POST", "path": "/mcp", + "scheme": "http", + "query_string": b"", + "root_path": "", + "server": ("localhost", 8000), "headers": [ (b"content-type", b"application/json"), + (b"host", b"localhost:8000"), ], } - receive = AsyncMock() + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + "more_body": False, + } + ) send = AsyncMock() user_auth = MagicMock() user_auth.user_id = "test-user-id" @@ -729,6 +746,11 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): "handle_request", new_callable=AsyncMock, ) as mock_handle_request, + patch.object( + session_manager_stateless, + "_server_instances", + {}, + ), ): await handle_streamable_http_mcp(scope, receive, send) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index d31dbdd4348..6804ea9f8fe 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -2,6 +2,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ +import json import os import re import sys @@ -20,7 +21,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ( ToolPermissionGuardrail, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) @@ -894,3 +895,153 @@ class TestToolPermissionGuardrailIntegration: is_allowed, rule_id, _ = guardrail._check_tool_permission("Read") assert is_allowed is False assert rule_id == "deny_read" + + +class TestToolPermissionGuardrailInMemoryUpdate: + """Regression: an in-memory params update (PUT /guardrails path) must rebuild + the compiled rule maps, not just self.rules, so the new rules are enforced + without reinitializing the guardrail.""" + + def _bash(self, command): + return ChatCompletionMessageToolCall( + function={"name": "Bash", "arguments": json.dumps({"command": command})}, + type="function", + ) + + def test_update_in_memory_recompiles_added_param_pattern(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[{"id": "native-bash", "tool_name": r"^Bash$", "decision": "allow"}], + default_action="deny", + on_disallowed_action="block", + ) + # No pattern yet: any Bash command is allowed. + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0] + is True + ) + + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="deny", + on_disallowed_action="block", + rules=[ + { + "id": "native-bash", + "tool_name": r"^Bash$", + "decision": "allow", + "allowed_param_patterns": { + "command": r"^(?!(echo blockme)$).*$" + }, + } + ], + ) + ) + + # The compiled map must be rebuilt, and enforcement must reflect it. + assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {}) + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0] + is False + ) + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo hello"))[0] is True + ) + + def test_update_in_memory_recompiles_tool_name_target(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[], + default_action="allow", + on_disallowed_action="block", + ) + # No rules: default_action allow lets Bash through. + assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is True + + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="allow", + on_disallowed_action="block", + rules=[{"id": "deny-bash", "tool_name": r"^Bash$", "decision": "deny"}], + ) + ) + + # A newly added deny rule (new id) must match -> its compiled target was rebuilt. + assert "deny-bash" in guardrail._compiled_rule_targets + assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is False + + def test_update_in_memory_preserves_rules_when_rules_absent(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[ + { + "id": "native-bash", + "tool_name": r"^Bash$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^(?!(echo blockme)$).*$"}, + } + ], + default_action="deny", + on_disallowed_action="block", + ) + assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {}) + + # A partial update that does not carry `rules` must NOT wipe the existing + # ruleset / compiled maps. + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="deny", + on_disallowed_action="block", + ) + ) + + assert len(guardrail.rules) == 1 + assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {}) + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0] + is False + ) + + def test_update_in_memory_rejects_invalid_regex_and_keeps_previous_rules(self): + """Regression: a live update whose rules contain an invalid regex must be + rejected atomically. The bad rule must not leak in as a compiled-target + wildcard (match-all), and the previously enforced ruleset must survive.""" + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[{"id": "deny-secret", "tool_name": r"^Secret$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="block", + ) + # Baseline: only "Secret" is denied; any other tool is allowed. + assert guardrail._check_tool_permission("Secret")[0] is False + assert guardrail._check_tool_permission("Other")[0] is True + + with pytest.raises(ValueError): + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="allow", + on_disallowed_action="block", + rules=[ + { + "id": "deny-secret", + "tool_name": r"^Secret$", + "decision": "deny", + }, + {"id": "bad", "tool_name": "[unclosed", "decision": "deny"}, + ], + ) + ) + + # The bad rule must not have leaked in, and the prior ruleset must hold. + assert "bad" not in guardrail._compiled_rule_targets + assert all(rule.id != "bad" for rule in guardrail.rules) + assert guardrail._check_tool_permission("Other")[0] is True + assert guardrail._check_tool_permission("Secret")[0] is False diff --git a/tests/test_litellm/proxy/test_caching_routes.py b/tests/test_litellm/proxy/test_caching_routes.py index 3e842d118dd..840ba054cc9 100644 --- a/tests/test_litellm/proxy/test_caching_routes.py +++ b/tests/test_litellm/proxy/test_caching_routes.py @@ -123,33 +123,100 @@ def test_cache_ping_failure(mock_redis_failure): assert "message" in error_details assert "litellm_cache_params" in error_details assert "health_check_cache_params" in error_details - assert "traceback" in error_details - # Verify specific error message - assert "invalid username-password pair" in error_details["message"] + # Verify generic static message (exception text must not leak to clients) + assert error_details["message"] == "Service Unhealthy" -def test_cache_ping_no_cache_initialized(): - """Test cache ping when no cache is initialized""" - # Set cache to None - original_cache = litellm.cache - litellm.cache = None - +def test_cache_ping_failure_does_not_expose_traceback(mock_redis_failure): + """CWE-209: Stack trace and exception text must not appear in the HTTP 503 response body.""" response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"}) assert response.status_code == 503 data = response.json() - print("response data=", json.dumps(data, indent=4)) - assert "error" in data - error = data["error"] + error = data.get("error", {}) + raw_body = json.dumps(data) - # Verify error contains all expected fields - assert "message" in error + # The word "traceback" (case-insensitive) must not appear anywhere in the response + assert ( + "traceback" not in raw_body.lower() + ), "CWE-209: Python traceback exposed in HTTP 503 response body" + # Internal frame paths should not leak either + assert ( + 'File "' not in raw_body + ), "CWE-209: Python stack frame paths exposed in HTTP 503 response body" + # Exception text (e.g. Redis hostnames/IPs) must not leak either + assert ( + "invalid username-password pair" not in raw_body + ), "CWE-209: Exception message text exposed in HTTP 503 response body" + + # The error message should be the safe static string error_details = json.loads(error["message"]) - assert "Cache not initialized. litellm.cache is None" in error_details["message"] + assert error_details["message"] == "Service Unhealthy" - # Restore original cache - litellm.cache = original_cache + +def test_cache_ping_no_cache_initialized(): + """Test cache ping when no cache is initialized returns 503 with ProxyException envelope. + + Verifies the exact response structure so that regressions in the error format + (e.g. message moving to a different field, or extra internal details leaking) + are caught immediately. + """ + original_cache = litellm.cache + litellm.cache = None + + try: + response = client.get( + "/cache/ping", headers={"Authorization": "Bearer sk-1234"} + ) + assert response.status_code == 503 + + data = response.json() + print("response data=", json.dumps(data, indent=4)) + # ProxyException is serialised as {"error": {"message": "...", "type": ..., ...}} + assert "error" in data + error_details = json.loads(data["error"]["message"]) + assert ( + error_details["message"] == "Cache not initialized. litellm.cache is None" + ) + finally: + litellm.cache = original_cache + + +def test_cache_ping_no_cache_does_not_expose_internals(): + """CWE-209: No-cache 503 must use the ProxyException envelope with no internal details. + + The null-cache path raises ProxyException directly (not HTTPException), so the + response is {"error": {"message": "...", ...}} — same envelope as other 503s from + this endpoint — with no tracebacks, source paths, or extra fields leaking. + """ + original_cache = litellm.cache + litellm.cache = None + + try: + response = client.get( + "/cache/ping", headers={"Authorization": "Bearer sk-1234"} + ) + assert response.status_code == 503 + + raw_body = response.text + # No Python traceback or source-file paths must appear in the response + assert "traceback" not in raw_body.lower(), ( + "CWE-209: Python traceback exposed in /cache/ping no-cache response" + ) + assert 'File "' not in raw_body, ( + "CWE-209: Python stack frame paths exposed in /cache/ping no-cache response" + ) + + data = response.json() + # Response must use the ProxyException envelope + assert "error" in data, f"Expected ProxyException envelope, got: {data}" + error_details = json.loads(data["error"]["message"]) + assert ( + error_details["message"] == "Cache not initialized. litellm.cache is None" + ) + finally: + litellm.cache = original_cache def test_cache_ping_health_check_includes_only_cache_attributes(mock_redis_success): diff --git a/tests/test_litellm/proxy/test_dynamic_mcp_route.py b/tests/test_litellm/proxy/test_dynamic_mcp_route.py index 2462aff2119..592cebd957c 100644 --- a/tests/test_litellm/proxy/test_dynamic_mcp_route.py +++ b/tests/test_litellm/proxy/test_dynamic_mcp_route.py @@ -486,3 +486,57 @@ async def test_dynamic_mcp_route_empty_access_group_returns_404(): await dynamic_mcp_route("empty_group", request) assert exc_info.value.status_code == 404 + + +# --------------------------------------------------------------------------- +# 6. Unexpected exception → 500 without leaking stack trace (CWE-209) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_unexpected_exception_returns_500_without_traceback(): + """CWE-209: an unexpected exception must return 500 with a generic message, + never leaking str(e) or a Python traceback to the caller.""" + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/boom/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock( + side_effect=RuntimeError("internal host: redis://10.0.0.1:6379") + ) + + with patch(_MCP_MANAGER, fake_mgr): + with pytest.raises(HTTPException) as exc_info: + await dynamic_mcp_route("boom", request) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Internal server error" + assert "10.0.0.1" not in str(exc_info.value.detail) + assert "traceback" not in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_toolset_mcp_route_unexpected_exception_returns_500_without_traceback(): + """CWE-209: toolset_mcp_route must return 500 with a generic message on + unexpected errors, never leaking exception text to the caller.""" + from litellm.proxy.proxy_server import toolset_mcp_route + + request = _make_request("/toolset/broken_toolset/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_toolset_by_name_cached = AsyncMock( + side_effect=RuntimeError("connection to db-host:5432 refused") + ) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_PRISMA, new=MagicMock()), + ): + with pytest.raises(HTTPException) as exc_info: + await toolset_mcp_route("broken_toolset", request) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Internal server error" + assert "db-host" not in str(exc_info.value.detail) + assert "traceback" not in str(exc_info.value.detail).lower() diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index b03579c2dbd..113e1bc0df8 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -849,12 +849,13 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( @pytest.mark.parametrize( - "model, model_info, expected_model_param", + "model, model_info, expected_model_param, expected_base_model_param", [ - ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro"), + ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None), ( "gemini/gemini-3.1-pro", {"base_model": "gemini-3.1-pro-preview"}, + "gemini-3.1-pro", "gemini-3.1-pro-preview", ), ], @@ -863,7 +864,13 @@ def test_completion_optional_params_base_model( model: str, model_info: dict | None, expected_model_param: str, + expected_base_model_param: str | None, ): + """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model`` + (an additive capability hint), without overwriting ``model`` with the label. + + Regression for #29618: overwriting ``model`` with a friendly ``base_model`` + label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``.""" with patch("litellm.main.get_optional_params") as mock_get_optional_params: mock_get_optional_params.return_value = MagicMock() @@ -881,10 +888,9 @@ def test_completion_optional_params_base_model( litellm.completion(**kwargs) assert mock_get_optional_params.called is True - get_optional_params_model_param = mock_get_optional_params.call_args.kwargs[ - "model" - ] - assert get_optional_params_model_param == expected_model_param + call_kwargs = mock_get_optional_params.call_args.kwargs + assert call_kwargs["model"] == expected_model_param + assert call_kwargs["base_model"] == expected_base_model_param @patch("litellm.completion_extras.responses_api_bridge.completion") diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2d75671f1cb..f2b9bef9230 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4144,3 +4144,51 @@ class TestValidateAndFixThinkingParam: validate_and_fix_thinking_param(thinking=thinking) assert "budgetTokens" in thinking assert "budget_tokens" not in thinking + + +class TestBedrockBaseModelLabelKeepsTools: + """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly + label must not silently drop ``tools``/``tool_choice`` under ``drop_params``.""" + + TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ] + + def test_base_model_label_keeps_tools_with_drop_params(self): + from litellm.utils import get_optional_params + + result = get_optional_params( + model="eu.anthropic.claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + base_model="claude-haiku-4-5", + tools=self.TOOLS, + tool_choice="auto", + drop_params=True, + ) + + assert "tools" in result + assert "tool_choice" in result + + def test_base_model_label_alone_drops_tools(self): + """Without the real model id the label resolves to no tool support, so passing + the label as ``model`` is exactly what dropped tools before the fix.""" + from litellm.utils import get_optional_params + + result = get_optional_params( + model="claude-haiku-4-5", + custom_llm_provider="bedrock", + tools=self.TOOLS, + tool_choice="auto", + drop_params=True, + ) + + assert "tools" not in result From 7edf3a9cb55548b143df1692f4ed7c4681d7fcf7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 4 Jun 2026 11:37:54 -0700 Subject: [PATCH 002/133] style(ui): run prettier --write across the dashboard (#29622) Formatting-only pass; no logic changes. Brings the UI into compliance with .prettierrc so the new format-check CI job passes --- ui/litellm-dashboard/e2e_tests/globalSetup.ts | 7 +- .../e2e_tests/helpers/navigation.ts | 4 +- .../e2e_tests/playwright.config.ts | 2 +- .../e2e_tests/tests/auth/logout.spec.ts | 9 +- .../tests/auth/proxyLogoutUrl.spec.ts | 16 +- .../tests/internal-user/internalUser.spec.ts | 12 +- .../internalUserWithTeams.spec.ts | 6 +- .../internal-viewer/internalViewer.spec.ts | 6 +- .../tests/login/internalUserIdentity.spec.ts | 4 +- .../e2e_tests/tests/login/login.spec.ts | 9 +- .../e2e_tests/tests/mcp/mcpServers.spec.ts | 6 +- .../e2e_tests/tests/modelHub/modelHub.spec.ts | 5 +- .../tests/modelsPage/addModel.spec.ts | 33 +- .../modelsPage/clearCustomPricing.spec.ts | 56 +- .../tests/navigation/sidebar.spec.ts | 10 +- .../e2e_tests/tests/proxy-admin/keys.spec.ts | 8 +- .../tests/proxy-admin/license.spec.ts | 9 +- .../e2e_tests/tests/proxy-admin/teams.spec.ts | 10 +- .../tests/settings/routerSettings.spec.ts | 9 +- .../tests/team-admin/teamAdmin.spec.ts | 12 +- ui/litellm-dashboard/knip.json | 13 +- .../src/app/(dashboard)/README.md | 8 +- .../app/(dashboard)/components/Sidebar2.tsx | 1 - .../components/SidebarProvider.tsx | 6 +- .../accessGroups/useAccessGroupDetails.ts | 20 +- .../hooks/accessGroups/useAccessGroups.ts | 14 +- .../accessGroups/useCreateAccessGroup.ts | 7 +- .../accessGroups/useDeleteAccessGroup.ts | 12 +- .../hooks/accessGroups/useEditAccessGroup.ts | 7 +- .../cloudzero/useCloudZeroCreate.test.ts | 35 +- .../cloudzero/useCloudZeroDryRun.test.ts | 35 +- .../cloudzero/useCloudZeroExport.test.ts | 35 +- .../hooks/common/queryKeysFactory.test.ts | 6 +- .../configOverrides/hashicorpVaultApi.ts | 17 +- .../hooks/guardrails/useGuardrails.test.ts | 8 +- .../hooks/guardrails/useRegisterGuardrail.ts | 7 +- .../useHealthReadinessDetails.ts | 13 +- .../hooks/keys/useKeyAliases.test.ts | 12 +- .../(dashboard)/hooks/keys/useKeyAliases.ts | 14 +- .../(dashboard)/hooks/keys/useKeys.test.ts | 15 +- .../src/app/(dashboard)/hooks/keys/useKeys.ts | 18 +- .../hooks/keys/useResetKeySpend.ts | 12 +- .../hooks/logDetails/useLogDetails.ts | 6 +- .../useMCPSemanticFilterSettings.ts | 4 +- .../useUpdateMCPSemanticFilterSettings.ts | 4 +- .../mcpServers/useMCPAccessGroups.test.ts | 2 +- .../hooks/mcpServers/useMCPServerHealth.ts | 36 +- .../hooks/mcpServers/useMCPServers.test.ts | 2 +- .../app/(dashboard)/hooks/models/useModels.ts | 27 +- .../hooks/projects/useCreateProject.test.ts | 4 +- .../hooks/projects/useCreateProject.ts | 12 +- .../hooks/projects/useDeleteProject.test.ts | 4 +- .../hooks/projects/useDeleteProject.ts | 12 +- .../hooks/projects/useProjectDetails.ts | 20 +- .../(dashboard)/hooks/projects/useProjects.ts | 11 +- .../hooks/projects/useUpdateProject.test.ts | 6 +- .../hooks/projects/useUpdateProject.ts | 13 +- .../storeModelInDB/useStoreModelInDB.test.ts | 9 +- .../hooks/storeModelInDB/useStoreModelInDB.ts | 8 +- .../useStoreRequestInSpendLogs.ts | 2 +- .../(dashboard)/hooks/teams/useTeams.test.ts | 25 +- .../app/(dashboard)/hooks/teams/useTeams.ts | 24 +- .../(dashboard)/hooks/useAuthorized.test.ts | 14 +- .../(dashboard)/hooks/users/useUsers.test.ts | 70 +- .../app/(dashboard)/hooks/users/useUsers.ts | 13 +- .../src/app/(dashboard)/layout.tsx | 8 +- .../ModelsAndEndpointsView.tsx | 6 +- .../components/AllModelsTab.test.tsx | 300 ++- .../components/AllModelsTab.tsx | 78 +- .../src/app/(dashboard)/playground/page.tsx | 74 +- .../src/app/login/LoginPage.test.tsx | 3 +- .../src/app/login/LoginPage.tsx | 19 +- .../src/app/mcp/oauth/callback/page.tsx | 12 +- .../src/app/model_hub_table/page.tsx | 4 +- .../onboarding/OnboardingErrorView.test.tsx | 4 +- .../src/app/onboarding/OnboardingForm.tsx | 10 +- .../onboarding/OnboardingFormBody.test.tsx | 8 +- .../src/app/onboarding/OnboardingFormBody.tsx | 36 +- .../src/app/onboarding/page.tsx | 6 +- ui/litellm-dashboard/src/app/page.tsx | 475 ++-- .../AIHub/AgentHubTableColumns.test.tsx | 10 +- .../AIHub/ClaudeCodeMarketplaceTab.tsx | 39 +- .../components/AIHub/ModelHubTable.test.tsx | 34 +- .../src/components/AIHub/ModelHubTable.tsx | 4 +- .../components/AIHub/SkillHubDashboard.tsx | 12 +- .../AIHub/marketplace_table_columns.tsx | 10 +- .../AccessGroupsDetailsPage.test.tsx | 108 +- .../AccessGroups/AccessGroupsDetailsPage.tsx | 79 +- .../AccessGroupsModal/AccessGroupBaseForm.tsx | 20 +- .../AccessGroupCreateModal.tsx | 11 +- .../AccessGroupEditModal.tsx | 17 +- .../AccessGroups/AccessGroupsPage.test.tsx | 68 +- .../AccessGroups/AccessGroupsPage.tsx | 104 +- .../src/components/AccessGroups/types.ts | 58 +- .../src/components/BulkEditUsers.test.tsx | 12 +- .../src/components/BulkEditUsers.tsx | 12 +- .../add_margin_form.test.tsx | 37 +- .../CostTrackingSettings/add_margin_form.tsx | 17 +- .../add_provider_form.test.tsx | 19 +- .../add_provider_form.tsx | 11 +- .../cost_tracking_settings.test.tsx | 34 +- .../cost_tracking_settings.tsx | 90 +- .../CostTrackingSettings/how_it_works.tsx | 202 +- .../components/CostTrackingSettings/index.ts | 9 +- .../pricing_calculator/index.test.tsx | 28 +- .../pricing_calculator/index.tsx | 31 +- .../multi_cost_results.test.tsx | 59 +- .../pricing_calculator/multi_cost_results.tsx | 73 +- .../multi_export_dropdown.test.tsx | 6 +- .../multi_export_dropdown.tsx | 8 +- .../multi_export_utils.test.ts | 66 +- .../pricing_calculator/multi_export_utils.ts | 19 +- .../pricing_calculator/types.ts | 1 - .../use_multi_cost_estimate.test.ts | 32 +- .../use_multi_cost_estimate.ts | 10 +- .../provider_discount_table.test.tsx | 40 +- .../provider_discount_table.tsx | 5 +- .../provider_display_helpers.ts | 7 +- .../provider_margin_table.test.tsx | 42 +- .../provider_margin_table.tsx | 10 +- .../components/CostTrackingSettings/types.ts | 1 - .../use_discount_config.test.ts | 12 +- .../use_discount_config.ts | 164 +- .../use_margin_config.test.ts | 12 +- .../CostTrackingSettings/use_margin_config.ts | 199 +- .../src/components/CreateUserButton.test.tsx | 107 +- .../src/components/CreateUserButton.tsx | 26 +- .../src/components/DebugWarningBanner.tsx | 7 +- .../src/components/DefaultUserSettings.tsx | 11 +- .../DeletedKeysPage/DeletedKeysPage.tsx | 6 +- .../DeletedKeysTable/DeletedKeysTable.tsx | 53 +- .../DeletedTeamsPage/DeletedTeamsPage.tsx | 12 +- .../DeletedTeamsTable.test.tsx | 8 +- .../DeletedTeamsTable/DeletedTeamsTable.tsx | 51 +- .../ExportFormatSelector.tsx | 1 - .../EntityUsageExport/ExportSummary.test.tsx | 20 +- .../EntityUsageExport/ExportSummary.tsx | 1 - .../ExportTypeSelector.test.tsx | 20 +- .../EntityUsageExport/UsageExportHeader.tsx | 2 +- .../src/components/EntityUsageExport/index.ts | 1 - .../EntityUsageExport/utils.test.ts | 1 - .../src/components/EntityUsageExport/utils.ts | 17 +- .../src/components/GuardrailSettingsView.tsx | 11 +- .../EvaluationSettingsModal.tsx | 10 +- .../GuardrailConfig.test.tsx | 4 +- .../GuardrailsMonitor/GuardrailConfig.tsx | 26 +- .../GuardrailsMonitor/GuardrailDetail.tsx | 38 +- .../GuardrailsMonitorView.test.tsx | 11 +- .../GuardrailsMonitorView.tsx | 11 +- .../GuardrailsMonitor/GuardrailsOverview.tsx | 71 +- .../GuardrailsMonitor/LogViewer.tsx | 37 +- .../GuardrailsMonitor/MetricCard.test.tsx | 12 +- .../GuardrailsMonitor/MetricCard.tsx | 12 +- .../GuardrailsMonitor/ScoreChart.tsx | 4 +- .../src/components/HelpLink.test.tsx | 14 +- .../src/components/HelpLink.tsx | 37 +- .../PaginatedKeyAliasSelect.test.tsx | 12 +- .../PaginatedKeyAliasSelect.tsx | 15 +- .../components/MemoryView/MemoryEditModal.tsx | 37 +- .../src/components/MemoryView/MemoryView.tsx | 156 +- .../ModelSelect/ModelSelect.test.tsx | 14 +- .../components/ModelSelect/ModelSelect.tsx | 98 +- .../PaginatedModelSelect.test.tsx | 16 +- .../PaginatedModelSelect.tsx | 23 +- .../Navbar/BlogDropdown/BlogDropdown.test.tsx | 4 +- .../WorkerDropdown/WorkerDropdown.test.tsx | 7 +- .../Navbar/WorkerDropdown/WorkerDropdown.tsx | 4 +- .../src/components/OldTeams.test.tsx | 22 +- .../src/components/OldTeams.tsx | 1531 ++++++------ .../Projects/ProjectDetailsPage.test.tsx | 3 +- .../Projects/ProjectKeysSection.test.tsx | 4 +- .../Projects/ProjectKeysTable.test.tsx | 13 +- .../ProjectModals/CreateProjectModal.tsx | 15 +- .../ProjectModals/EditProjectModal.test.tsx | 24 +- .../ProjectModals/EditProjectModal.tsx | 21 +- .../ProjectModals/ProjectBaseForm.tsx | 131 +- .../ProjectModals/projectFormUtils.ts | 7 +- .../components/Projects/ProjectsPage.test.tsx | 12 +- .../src/components/Projects/ProjectsPage.tsx | 44 +- .../SearchTools/CreateSearchTools.tsx | 10 +- .../SearchTools/SearchConnectionTest.tsx | 20 +- .../SearchTools/SearchToolColumn.tsx | 180 +- .../SearchTools/SearchToolTester.test.tsx | 5 +- .../SearchTools/SearchToolTester.tsx | 106 +- .../SearchTools/SearchToolView.test.tsx | 38 +- .../components/SearchTools/SearchToolView.tsx | 33 +- .../components/SearchTools/SearchTools.tsx | 38 +- .../src/components/SearchTools/index.tsx | 11 +- .../src/components/SearchTools/types.tsx | 1 - .../EditHashicorpVaultModal.tsx | 23 +- .../HashicorpVault/HashicorpVault.tsx | 52 +- .../HashicorpVaultEmptyPlaceholder.test.tsx | 4 +- .../HashicorpVaultEmptyPlaceholder.tsx | 3 +- .../AdminSettings/HashicorpVault/constants.ts | 6 +- .../LoggingSettings/LoggingSettings.tsx | 16 +- .../MCPSemanticFilterSettings.test.tsx | 22 +- .../MCPSemanticFilterSettings.tsx | 23 +- .../MCPSemanticFilterTestPanel.test.tsx | 31 +- .../MCPSemanticFilterTestPanel.tsx | 188 +- .../semanticFilterTestUtils.test.ts | 6 +- .../semanticFilterTestUtils.ts | 14 +- .../AdminSettings/SSOSettings/SSOSettings.tsx | 37 +- .../PageVisibilitySettings.test.tsx | 30 +- .../AdminSettings/UISettings/UISettings.tsx | 8 +- .../RouterSettings/Fallbacks/AddFallbacks.tsx | 17 +- .../Fallbacks/AddFallbacksModal.test.tsx | 4 +- .../Fallbacks/AddFallbacksModal.tsx | 6 +- .../Fallbacks/FallbackGroupConfig.tsx | 31 +- .../Fallbacks/FallbackSelectionForm.test.tsx | 78 +- .../Fallbacks/FallbackSelectionForm.tsx | 10 +- .../Fallbacks/Fallbacks.test.tsx | 17 +- .../RouterSettings/Fallbacks/Fallbacks.tsx | 29 +- .../src/components/TeamSSOSettings.test.tsx | 13 +- .../src/components/TeamSSOSettings.tsx | 8 +- .../src/components/ToolDetail.tsx | 59 +- .../src/components/ToolPolicies.tsx | 19 +- .../ToolPolicies/PolicySelect.test.tsx | 43 +- .../components/ToolPolicies/PolicySelect.tsx | 3 +- .../src/components/ToolPoliciesView.tsx | 16 +- .../src/components/UsageIndicator.tsx | 50 +- .../components/EndpointUsageBarChart.test.tsx | 4 +- .../components/EntityUsage/EntityUsage.tsx | 5 +- .../EntityUsage/SpendByProvider.test.tsx | 8 +- .../EntityUsage/SpendByProvider.tsx | 8 +- .../components/KeyModelUsageView.tsx | 6 +- .../components/UsageAIChatPanel.test.tsx | 9 +- .../UsagePage/components/UsageAIChatPanel.tsx | 66 +- .../components/UsagePageView.test.tsx | 42 +- .../UsageViewSelect/UsageViewSelect.tsx | 11 +- .../hooks/usePaginatedDailyActivity.ts | 13 +- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 40 +- .../VirtualKeysPage/VirtualKeysTable.tsx | 904 ++++--- .../src/components/WebRTCTester.jsx | 295 ++- .../src/components/activity_metrics.test.tsx | 12 +- .../src/components/activity_metrics.tsx | 18 +- .../add_model/AddModelForm.test.tsx | 13 +- .../src/components/add_model/AddModelForm.tsx | 4 +- .../add_model/ComplexityRouterConfig.test.tsx | 40 +- .../add_model/RouterConfigBuilder.test.tsx | 11 +- .../add_model/RouterConfigBuilder.tsx | 22 +- .../add_model/add_auto_router_tab.tsx | 52 +- .../handle_add_auto_router_submit.tsx | 5 +- .../add_model/handle_add_model_submit.tsx | 18 +- .../add_model/model_connection_test.tsx | 6 +- .../src/components/add_pass_through.tsx | 41 +- .../agent_management/AgentSelector.test.tsx | 13 +- .../agent_management/AgentSelector.tsx | 21 +- .../src/components/agents.test.tsx | 4 +- .../src/components/agents.tsx | 26 +- .../src/components/agents/add_agent_form.tsx | 152 +- .../src/components/agents/agent_card.test.tsx | 10 +- .../src/components/agents/agent_card.tsx | 25 +- .../agents/agent_card_discovery.test.tsx | 95 +- .../agents/agent_card_discovery.tsx | 106 +- .../agents/agent_card_grid.test.tsx | 20 +- .../src/components/agents/agent_cost_view.tsx | 13 +- .../agents/agent_discovery_utils.test.ts | 22 +- .../agents/agent_discovery_utils.ts | 24 +- .../components/agents/agent_form_fields.tsx | 404 ++-- .../src/components/agents/agent_info.tsx | 157 +- .../src/components/agents/agent_table.tsx | 32 +- .../src/components/agents/agent_type_utils.ts | 8 +- .../components/agents/cost_config_fields.tsx | 8 +- .../agents/dynamic_agent_form_fields.tsx | 30 +- .../components/alerting/alerting_settings.tsx | 1 - .../src/components/atoms/Tooltip.test.tsx | 6 +- .../components/budgets/budget_panel.test.tsx | 7 +- .../src/components/budgets/budget_panel.tsx | 5 +- .../components/budgets/edit_budget_modal.tsx | 21 +- .../components/bulk_create_users_button.tsx | 30 +- .../cache_settings/CacheFieldGroup.tsx | 11 +- .../src/components/chat/ChatMessages.tsx | 155 +- .../src/components/chat/ChatPage.tsx | 723 ++++-- .../src/components/chat/ConversationList.tsx | 51 +- .../src/components/chat/MCPAppsPanel.tsx | 307 ++- .../src/components/chat/MCPConnectPicker.tsx | 21 +- .../src/components/chat/MCPCredentialsTab.tsx | 35 +- .../src/components/chat/useChatHistory.ts | 142 +- .../src/components/claude_code_plugins.tsx | 21 +- .../MakeSkillPublicForm.tsx | 35 +- .../add_plugin_form.test.tsx | 20 +- .../claude_code_plugins/add_plugin_form.tsx | 82 +- .../claude_code_plugins/helpers.test.ts | 6 +- .../components/claude_code_plugins/helpers.ts | 48 +- .../claude_code_plugins/plugin_info.tsx | 66 +- .../claude_code_plugins/plugin_table.tsx | 124 +- .../claude_code_plugins/skill_detail.tsx | 74 +- .../components/claude_code_plugins/types.ts | 14 +- .../common_components/AccessGroupSelector.tsx | 14 +- .../DefaultProxyAdminTag.tsx | 4 +- .../DeleteResourceModal.test.tsx | 12 +- .../common_components/FilterTeamDropdown.tsx | 7 +- .../KeyLifecycleSettings.test.tsx | 72 +- .../KeyLifecycleSettings.tsx | 4 +- .../common_components/LabeledField.test.tsx | 16 +- .../common_components/MemberTable.tsx | 6 +- .../components/common_components/NewBadge.tsx | 8 +- .../PassThroughGuardrailsSection.tsx | 69 +- .../PassThroughRoutesSelector.tsx | 15 +- .../common_components/ProjectDropdown.tsx | 13 +- .../RateLimitTypeFormItem.test.tsx | 14 +- .../RouterSettingsAccordion.tsx | 16 +- .../TableHeaderSortDropdown.tsx | 5 +- .../common_components/simple_table.tsx | 1 - .../common_components/team_dropdown.tsx | 14 +- .../common_components/team_multi_select.tsx | 14 +- .../src/components/general_settings.tsx | 20 +- .../src/components/guardrails.tsx | 7 +- .../guardrails/GuardrailTestPanel.test.tsx | 3 +- .../guardrails/GuardrailTestPanel.tsx | 20 +- .../guardrails/GuardrailTestPlayground.tsx | 34 +- .../guardrails/GuardrailTestResults.test.tsx | 1 - .../guardrails/GuardrailTestResults.tsx | 22 +- .../guardrails/TeamGuardrailsTab.tsx | 290 +-- .../guardrails/add_guardrail_form.tsx | 51 +- .../content_filter/CategoryTable.tsx | 36 +- .../CompetitorIntentConfiguration.tsx | 55 +- .../ContentCategoryConfiguration.tsx | 76 +- .../ContentFilterConfiguration.tsx | 70 +- .../content_filter/ContentFilterDisplay.tsx | 9 +- .../ContentFilterManager.test.tsx | 74 +- .../content_filter/ContentFilterManager.tsx | 29 +- .../CustomPatternModal.test.tsx | 3 +- .../content_filter/CustomPatternModal.tsx | 19 +- .../content_filter/KeywordModal.tsx | 19 +- .../content_filter/KeywordTable.tsx | 31 +- .../content_filter/PatternModal.test.tsx | 9 +- .../content_filter/PatternModal.tsx | 21 +- .../content_filter/PatternTable.tsx | 44 +- .../custom_code/CustomCodeModal.tsx | 216 +- .../guardrails/edit_guardrail_form.tsx | 12 +- .../guardrails/guardrail_garden.tsx | 33 +- .../guardrails/guardrail_garden_card.tsx | 5 +- .../guardrails/guardrail_garden_data.ts | 18 +- .../guardrails/guardrail_garden_detail.tsx | 40 +- .../guardrails/guardrail_info.test.tsx | 21 +- .../components/guardrails/guardrail_info.tsx | 111 +- .../guardrail_info_helpers.test.tsx | 16 +- .../guardrails/guardrail_info_helpers.tsx | 2 +- .../guardrails/guardrail_optional_params.tsx | 7 +- .../guardrails/guardrail_provider_fields.tsx | 10 +- .../guardrails/llm_judge/LLMJudgeFields.tsx | 24 +- .../ToolPermissionRulesEditor.test.tsx | 4 +- .../ToolPermissionRulesEditor.tsx | 38 +- .../key_team_helpers/BudgetWindowsEditor.tsx | 27 +- .../key_team_helpers/filter_helpers.test.ts | 16 +- .../key_team_helpers/filter_logic.test.tsx | 67 +- .../src/components/leftnav.tsx | 67 +- .../src/components/logging_settings_view.tsx | 16 +- .../src/components/mcp_hub_table_columns.tsx | 17 +- .../MCPToolPermissions.tsx | 8 +- .../mcp_tools/ByokCredentialModal.tsx | 33 +- .../components/mcp_tools/MCPLogoSelector.tsx | 11 +- .../mcp_tools/MCPNetworkSettings.tsx | 17 +- .../MCPPermissionManagement.test.tsx | 63 +- .../mcp_tools/MCPPermissionManagement.tsx | 30 +- .../mcp_tools/MCPStandardsSettings.test.tsx | 10 +- .../mcp_tools/MCPSubmissionsTab.tsx | 73 +- .../mcp_tools/MCPToolArgumentsForm.tsx | 42 +- .../components/mcp_tools/MCPToolsetsTab.tsx | 113 +- .../mcp_tools/McpCrudPermissionPanel.tsx | 90 +- .../components/mcp_tools/OAuthFormFields.tsx | 72 +- .../mcp_tools/OpenAPIFormSection.tsx | 6 +- .../mcp_tools/OpenAPIQuickPicker.tsx | 13 +- .../components/mcp_tools/ToolTestPanel.tsx | 89 +- .../mcp_tools/create_mcp_server.tsx | 57 +- .../mcp_tools/mcp_connection_status.test.tsx | 30 +- .../mcp_tools/mcp_connection_status.tsx | 27 +- .../components/mcp_tools/mcp_discovery.tsx | 24 +- .../mcp_tools/mcp_server_columns.tsx | 12 +- .../components/mcp_tools/mcp_server_edit.tsx | 11 +- .../components/mcp_tools/mcp_server_view.tsx | 90 +- .../components/mcp_tools/mcp_servers.test.tsx | 10 +- .../src/components/mcp_tools/mcp_servers.tsx | 117 +- .../src/components/mcp_tools/mcp_tools.tsx | 305 ++- .../src/components/mcp_tools/utils.test.tsx | 8 +- .../components/model_add/credentials.test.tsx | 4 +- .../src/components/model_add/credentials.tsx | 4 +- .../ModelSettingsModal.test.tsx | 18 +- .../ModelSettingsModal/ModelSettingsModal.tsx | 6 +- .../model_dashboard/all_models_table.tsx | 32 +- .../src/components/model_info_view.test.tsx | 4 +- .../src/components/model_info_view.tsx | 51 +- .../molecules/models/columns.test.tsx | 40 +- .../components/molecules/models/columns.tsx | 769 +++--- .../src/components/networking.tsx | 295 +-- .../components/object_permissions_view.tsx | 6 +- .../organisms/RegenerateKeyModal.tsx | 8 +- .../organisms/create_key_button.test.tsx | 87 +- .../organisms/create_key_button.tsx | 14 +- .../organization/organization_view.tsx | 37 +- .../src/components/page_utils.test.ts | 109 +- .../src/components/page_utils.ts | 6 +- .../src/components/pass_through_info.tsx | 20 +- .../src/components/pass_through_settings.tsx | 8 +- .../permissions/AgentPermissions.tsx | 17 +- .../permissions/MCPServerPermissions.test.tsx | 25 +- .../permissions/MCPServerPermissions.tsx | 111 +- .../playground/chat_ui/A2AMetrics.tsx | 12 +- .../chat_ui/AdditionalModelSettings.test.tsx | 12 +- .../chat_ui/AdditionalModelSettings.tsx | 9 +- .../playground/chat_ui/AgentBuilderView.tsx | 77 +- .../chat_ui/ChatMessageBubble.test.tsx | 48 +- .../playground/chat_ui/ChatMessageBubble.tsx | 8 +- .../playground/chat_ui/ChatUI.test.tsx | 8 +- .../components/playground/chat_ui/ChatUI.tsx | 2136 ++++++++--------- .../chat_ui/CodeInterpreterOutput.tsx | 54 +- .../chat_ui/CodeInterpreterTool.tsx | 2 +- .../playground/chat_ui/CodeSnippets.tsx | 2 +- .../chat_ui/FilePreviewCard.test.tsx | 32 +- .../playground/chat_ui/FilePreviewCard.tsx | 4 +- .../playground/chat_ui/RealtimePlayground.tsx | 59 +- .../chat_ui/SearchResultsDisplay.tsx | 5 +- .../playground/chat_ui/useChatHistory.test.ts | 8 +- .../playground/chat_ui/useChatHistory.ts | 8 +- .../playground/chat_ui/useCodeInterpreter.ts | 2 +- .../playground/compareUI/CompareUI.test.tsx | 2 +- .../playground/compareUI/CompareUI.tsx | 15 +- .../compareUI/components/ComparisonPanel.tsx | 4 +- .../components/MessageInput.test.tsx | 8 +- .../compareUI/components/UnifiedSelector.tsx | 13 +- .../playground/compareUI/endpoint_config.ts | 16 +- .../playground/complianceUI/ComplianceUI.tsx | 1780 +++++++------- .../playground/llm_calls/a2a_send_message.tsx | 8 +- .../playground/llm_calls/chat_completion.tsx | 23 +- .../llm_calls/code_interpreter_handler.ts | 17 +- .../playground/llm_calls/fetch_agents.tsx | 8 +- .../playground/llm_calls/interactions_api.tsx | 4 +- .../policies/PolicySelector.test.tsx | 10 +- .../components/policies/PolicySelector.tsx | 4 +- .../policies/add_attachment_form.test.tsx | 5 +- .../policies/add_attachment_form.tsx | 73 +- .../components/policies/add_policy_form.tsx | 65 +- .../policies/ai_suggestion_modal.tsx | 686 +++--- .../policies/attachment_table.test.tsx | 37 +- .../policies/build_attachment_data.test.ts | 2 +- .../policies/build_attachment_data.ts | 2 +- .../policies/guardrail_selection_modal.tsx | 61 +- .../policies/impact_popover.test.tsx | 14 +- .../components/policies/impact_popover.tsx | 21 +- .../policies/impact_preview_alert.tsx | 38 +- .../src/components/policies/index.test.tsx | 13 +- .../src/components/policies/index.tsx | 58 +- .../policies/pipeline_flow_builder.tsx | 262 +- .../src/components/policies/policy_info.tsx | 25 +- .../components/policies/policy_table.test.tsx | 34 +- .../src/components/policies/policy_table.tsx | 11 +- .../policies/policy_templates.test.tsx | 20 +- .../components/policies/policy_templates.tsx | 62 +- .../components/policies/policy_test_panel.tsx | 62 +- .../policies/template_parameter_modal.tsx | 48 +- .../src/components/price_data_reload.tsx | 10 +- .../DeveloperMessageCard.tsx | 17 +- .../prompt_editor_view/DotpromptViewTab.tsx | 13 +- .../prompt_editor_view/ModelConfigCard.tsx | 13 +- .../prompt_editor_view/PromptCodeSnippets.tsx | 67 +- .../prompt_editor_view/PromptEditorHeader.tsx | 20 +- .../prompt_editor_view/PromptMessagesCard.tsx | 13 +- .../prompt_editor_view/PublishModal.tsx | 3 +- .../prompts/prompt_editor_view/ToolsCard.tsx | 23 +- .../VersionHistorySidePanel.tsx | 28 +- .../conversation_panel/EmptyState.tsx | 1 - .../conversation_panel/MessageBubble.tsx | 24 +- .../conversation_panel/MessageInput.tsx | 6 +- .../conversation_panel/MessageList.tsx | 8 +- .../conversation_panel/VariableInput.tsx | 15 +- .../conversation_panel/VariableWarning.tsx | 12 +- .../conversation_panel/index.tsx | 5 +- .../conversation_panel/types.ts | 1 - .../prompts/prompt_editor_view/index.tsx | 25 +- .../prompts/prompt_editor_view/types.ts | 1 - .../prompts/prompt_editor_view/utils.ts | 8 +- .../src/components/prompts/prompt_info.tsx | 107 +- .../src/components/prompts/prompt_table.tsx | 37 +- .../src/components/prompts/prompt_utils.tsx | 19 +- .../src/components/prompts/tool_modal.tsx | 7 +- .../components/prompts/variable_textarea.tsx | 18 +- .../src/components/public_model_hub.tsx | 6 +- .../src/components/query_param_input.tsx | 2 +- .../src/components/route_preview.tsx | 8 +- .../LatencyBasedConfiguration.test.tsx | 10 +- .../LatencyBasedConfiguration.tsx | 13 +- .../ReliabilityRetriesSection.test.tsx | 20 +- .../ReliabilityRetriesSection.tsx | 1 - .../RouterSettingsForm.test.tsx | 20 +- .../RoutingStrategySelector.test.tsx | 12 +- .../RoutingStrategySelector.tsx | 12 +- .../TagFilteringToggle.test.tsx | 44 +- .../router_settings/TagFilteringToggle.tsx | 13 +- .../components/router_settings/index.test.tsx | 32 +- .../src/components/router_settings/index.tsx | 4 +- .../routing_groups/RoutingGroupModal.tsx | 22 +- .../routing_groups/RoutingGroupsTable.tsx | 16 +- .../src/components/routing_groups/index.tsx | 24 +- .../src/components/routing_groups/types.ts | 6 +- .../components/shared/CreatedKeyDisplay.tsx | 8 +- .../components/skill_hub_table_columns.tsx | 18 +- .../survey/ClaudeCodeModal.test.tsx | 26 +- .../src/components/survey/ClaudeCodeModal.tsx | 14 +- .../survey/ClaudeCodePrompt.test.tsx | 20 +- .../components/survey/ClaudeCodePrompt.tsx | 3 +- .../src/components/survey/NudgePrompt.tsx | 19 +- .../components/survey/SurveyModal.test.tsx | 80 +- .../src/components/survey/SurveyModal.tsx | 24 +- .../components/survey/SurveyPrompt.test.tsx | 20 +- .../src/components/survey/SurveyPrompt.tsx | 1 - .../src/components/survey/index.tsx | 1 - .../src/components/team/EditMembership.tsx | 8 +- .../src/components/team/MyUserTab.tsx | 23 +- .../src/components/team/TeamInfo.test.tsx | 39 +- .../src/components/team/TeamInfo.tsx | 183 +- .../src/components/team/TeamMemberTab.tsx | 16 +- .../team/TeamVirtualKeysTable.test.tsx | 48 +- .../components/team/TeamVirtualKeysTable.tsx | 55 +- .../src/components/team/available_teams.tsx | 6 +- .../team/permission_definitions.tsx | 16 +- .../team/tabVisibilityUtils.test.ts | 6 +- .../src/components/team/tabVisibilityUtils.ts | 11 +- .../src/components/team/useMyTeamMember.ts | 15 +- .../components/templates/KeyInfoHeader.tsx | 28 +- .../key_info_view.budget_display.test.tsx | 4 +- .../templates/key_info_view.test.tsx | 102 +- .../components/templates/key_info_view.tsx | 74 +- .../components/ui/AntDLoadingSpinner.test.tsx | 7 +- .../src/components/ui_theme_settings.tsx | 71 +- .../src/components/user_dashboard.test.tsx | 12 +- .../src/components/user_edit_view.tsx | 12 +- .../CreateVectorStore.test.tsx | 4 +- .../CreateVectorStore.tsx | 28 +- .../S3VectorsConfig.test.tsx | 2 +- .../S3VectorsConfig.tsx | 14 +- .../TestVectorStoreTab.tsx | 6 +- .../VectorStoreForm.tsx | 4 +- .../VectorStoreTable.tsx | 19 +- .../vector_store_management/index.tsx | 13 +- .../AuditLogDrawer/AuditLogDrawer.tsx | 36 +- .../view_logs/CostBreakdownViewer.test.tsx | 42 +- .../view_logs/CostBreakdownViewer.tsx | 289 +-- .../components/view_logs/ErrorViewer.test.tsx | 4 +- .../view_logs/EvalViewer/EvalViewer.tsx | 31 +- .../GuardrailViewer/CompliancePanel.tsx | 22 +- .../GuardrailViewer/ContentFilterDetails.tsx | 6 +- .../GuardrailViewer/GuardrailViewer.tsx | 104 +- .../CollapsibleMessage.test.tsx | 20 +- .../LogDetailsDrawer/CollapsibleMessage.tsx | 44 +- .../LogDetailsDrawer/DrawerHeader.tsx | 10 +- .../LogDetailsDrawer/HistoryTree.test.tsx | 12 +- .../LogDetailsDrawer/HistoryTree.tsx | 38 +- .../LogDetailsDrawer/InputCard.test.tsx | 2 +- .../view_logs/LogDetailsDrawer/InputCard.tsx | 32 +- .../LogDetailsDrawer/LogDetailContent.tsx | 76 +- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 123 +- .../LogDetailsDrawer/OutputCard.test.tsx | 2 +- .../view_logs/LogDetailsDrawer/OutputCard.tsx | 49 +- .../LogDetailsDrawer/PrettyMessagesView.tsx | 14 +- .../RealtimePrettyView.test.tsx | 26 +- .../LogDetailsDrawer/RealtimePrettyView.tsx | 266 +- .../LogDetailsDrawer/SectionHeader.tsx | 71 +- .../SimpleMessageBlock.test.tsx | 24 +- .../LogDetailsDrawer/SimpleMessageBlock.tsx | 37 +- .../SimpleToolCallBlock.test.tsx | 20 +- .../LogDetailsDrawer/SimpleToolCallBlock.tsx | 28 +- .../LogDetailsDrawer/prettyMessagesTypes.ts | 2 +- .../LogDetailsDrawer/prettyMessagesUtils.ts | 84 +- .../components/view_logs/LogsTableToolbar.tsx | 13 +- .../view_logs/RequestResponsePanel.test.tsx | 8 +- .../ToolsSection/FormattedToolView.tsx | 28 +- .../ToolsSection/ToolExpandedContent.tsx | 12 +- .../view_logs/ToolsSection/ToolItem.tsx | 4 +- .../view_logs/ToolsSection/ToolsSection.tsx | 2 +- .../view_logs/ToolsSection/utils.ts | 33 +- .../src/components/view_logs/TypeBadges.tsx | 36 +- .../view_logs/VectorStoreViewer.tsx | 174 +- .../src/components/view_logs/audit_logs.tsx | 96 +- .../src/components/view_logs/columns.tsx | 6 +- .../src/components/view_logs/table.tsx | 25 +- .../src/components/view_logs/utils.ts | 9 +- .../src/components/view_users.tsx | 27 +- .../src/components/view_users/columns.tsx | 3 +- .../src/components/view_users/table.test.tsx | 8 +- .../view_users/user_info_view.test.tsx | 9 +- .../components/view_users/user_info_view.tsx | 137 +- .../src/components/workflow_runs/index.tsx | 140 +- .../src/contexts/ThemeContext.tsx | 4 +- .../src/data/canadianPiiCompliancePrompts.ts | 102 +- .../src/data/claimsCompliancePrompts.ts | 104 +- .../data/codeExecutionCompliancePrompts.ts | 220 +- .../src/data/compliancePrompts.ts | 65 +- .../src/data/financialCompliancePrompts.ts | 624 +++-- .../src/data/insultsCompliancePrompts.ts | 900 ++++--- .../useDeletePolicyAttachment.test.tsx | 4 +- .../policies/useDeletePolicyAttachment.ts | 6 +- .../src/hooks/useMcpOAuthFlow.tsx | 18 +- .../src/hooks/useTestMCPConnection.tsx | 57 +- .../src/hooks/useToolsOAuthFlow.tsx | 10 +- .../src/hooks/useUserMcpOAuthFlow.tsx | 4 +- ui/litellm-dashboard/src/hooks/useWorker.ts | 3 +- .../src/utils/cookieUtils.test.ts | 24 +- ui/litellm-dashboard/src/utils/errorUtils.ts | 16 +- .../src/utils/mcpTokenStore.test.ts | 8 +- .../src/utils/mcpTokenStore.ts | 11 +- .../src/utils/proxyUtils.test.ts | 2 +- .../src/utils/returnUrlUtils.test.ts | 8 +- ui/litellm-dashboard/src/utils/roles.test.ts | 4 +- ui/litellm-dashboard/src/utils/roles.ts | 6 +- .../src/utils/secureStorage.ts | 9 +- .../tests/CreateKeyPage.expiredToken.test.tsx | 2 +- ui/litellm-dashboard/tsconfig.json | 24 +- 608 files changed, 13772 insertions(+), 16368 deletions(-) diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts index 6ff5522244a..8f80f57bd78 100644 --- a/ui/litellm-dashboard/e2e_tests/globalSetup.ts +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -14,10 +14,9 @@ async function globalSetup() { await page.getByPlaceholder("Enter your username").fill(email); await page.getByPlaceholder("Enter your password").fill(password); await page.getByRole("button", { name: "Login", exact: true }).click(); - await page.waitForURL( - (url) => url.pathname.startsWith("/ui") && !url.pathname.includes("/login"), - { timeout: 30_000 }, - ); + await page.waitForURL((url) => url.pathname.startsWith("/ui") && !url.pathname.includes("/login"), { + timeout: 30_000, + }); await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); // Dismiss feedback popup if present const dismiss = page.getByText("Don't ask me again"); diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts index 556e964842a..6ca18890f7a 100644 --- a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts +++ b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts @@ -20,7 +20,9 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise if (await dismissButton.isVisible({ timeout: 1_500 }).catch(() => false)) { await dismissButton.click(); // Wait for the popup to disappear - await expect(dismissButton).not.toBeVisible({ timeout: 2_000 }).catch(() => {}); + await expect(dismissButton) + .not.toBeVisible({ timeout: 2_000 }) + .catch(() => {}); } } diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts index 6964fe52a14..8d586ce9503 100644 --- a/ui/litellm-dashboard/e2e_tests/playwright.config.ts +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -31,7 +31,7 @@ export default defineConfig({ /* Slow down actions when SLOWMO= is set, useful for headed local debugging */ launchOptions: { - slowMo: process.env.SLOWMO ? (parseInt(process.env.SLOWMO, 10) || 0) : 0, + slowMo: process.env.SLOWMO ? parseInt(process.env.SLOWMO, 10) || 0 : 0, }, }, diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts index fefadf27548..d8644babfe3 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts @@ -13,9 +13,12 @@ test.describe("Logout", () => { // is declared with trigger={["click"]}, so a plain click opens the popup. await page.getByRole("button", { name: /Account menu/i }).click(); - const popup = page.locator(".ant-dropdown:visible").filter({ - has: page.locator(".bg-white.rounded-lg.shadow-lg"), - }).first(); + const popup = page + .locator(".ant-dropdown:visible") + .filter({ + has: page.locator(".bg-white.rounded-lg.shadow-lg"), + }) + .first(); await expect(popup).toBeVisible({ timeout: 5_000 }); // Click Logout — the handler clears the auth cookie and navigates via diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts index 4a233ed1bb1..6358fcf438e 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts @@ -38,10 +38,9 @@ test.describe("PROXY_LOGOUT_URL redirect", () => { // fetch (/sso/get/ui_settings) resolves. Clicking Logout before that lands // runs `window.location.href = ""` — a same-origin reload, not a redirect — // so gate the click on the settings response, not just on first paint. - const settingsLoaded = page.waitForResponse( - (r) => r.url().includes("/sso/get/ui_settings") && r.ok(), - { timeout: 30_000 }, - ); + const settingsLoaded = page.waitForResponse((r) => r.url().includes("/sso/get/ui_settings") && r.ok(), { + timeout: 30_000, + }); await page.goto("/ui"); await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); await settingsLoaded; @@ -59,10 +58,7 @@ test.describe("PROXY_LOGOUT_URL redirect", () => { // handleLogout clears cookies/local storage, then assigns window.location.href. // Arm the navigation wait before the click so we never miss the redirect. - await Promise.all([ - page.waitForURL((url) => url.origin === target.origin, { timeout: 15_000 }), - logout.click(), - ]); + await Promise.all([page.waitForURL((url) => url.origin === target.origin, { timeout: 15_000 }), logout.click()]); // The browser landed on exactly the configured logout URL. Compare normalized // hrefs (both sides through URL()) so trailing-slash / default-port rewrites the @@ -74,9 +70,7 @@ test.describe("PROXY_LOGOUT_URL redirect", () => { // ...and the client-side session cookie is gone (clearTokenCookies ran before // the redirect). HttpOnly cookies set server-side can't be cleared from JS, // so scope the check to the JS-managed token the UI is responsible for. - const clientTokensAfter = (await page.context().cookies()).filter( - (c) => c.name === "token" && !c.httpOnly, - ); + const clientTokensAfter = (await page.context().cookies()).filter((c) => c.name === "token" && !c.httpOnly); expect(clientTokensAfter, "client token cookie should be cleared on logout").toHaveLength(0); }); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts index c706ae0aefc..07a75dc007d 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts @@ -22,9 +22,9 @@ test.describe("Internal User", () => { const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" }); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await expect( - page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first(), - ).toBeVisible({ timeout: 5_000 }); + await expect(page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({ + timeout: 5_000, + }); }); test("Team info page omits the Settings tab for non-admin members", async ({ page }) => { @@ -44,9 +44,9 @@ test.describe("Internal User", () => { // Anchor on the user's own seeded key so the absence check below cannot // pass vacuously against an empty table. - await expect( - page.locator("table tbody").getByText(E2E_INTERNAL_USER_KEY_ALIAS).first(), - ).toBeVisible({ timeout: 10_000 }); + await expect(page.locator("table tbody").getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()).toBeVisible({ + timeout: 10_000, + }); // The litellm-dashboard team is the proxy's internal bookkeeping team — // its keys must never leak into an internal user's Virtual Keys table. diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts index f6b60f411b4..7d5058a8140 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts @@ -1,9 +1,5 @@ import { test, expect } from "@playwright/test"; -import { - INTERNAL_USER_STORAGE_PATH, - E2E_TEAM_CRUD_ALIAS, - E2E_TEAM_ORG_ALIAS, -} from "../../constants"; +import { INTERNAL_USER_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_ORG_ALIAS } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts index f5ab3c00503..4de86c46398 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts @@ -1,9 +1,5 @@ import { test, expect } from "@playwright/test"; -import { - E2E_TEAM_CRUD_ID, - E2E_VIEWER_KEY_ALIAS, - INTERNAL_VIEWER_STORAGE_PATH, -} from "../../constants"; +import { E2E_TEAM_CRUD_ID, E2E_VIEWER_KEY_ALIAS, INTERNAL_VIEWER_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts index cbe95276929..6008049a2aa 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts @@ -22,9 +22,7 @@ test.describe("Navbar identity scoping", () => { await expect(accountButton).toHaveAttribute("aria-label", /Internal User/, { timeout: 5_000 }); await expect(accountButton).toHaveAttribute( "aria-label", - new RegExp( - `signed in as (${escapeRegExp(E2E_INTERNAL_USER_EMAIL)}|${escapeRegExp(E2E_INTERNAL_USER_ID)})`, - ), + new RegExp(`signed in as (${escapeRegExp(E2E_INTERNAL_USER_EMAIL)}|${escapeRegExp(E2E_INTERNAL_USER_ID)})`), { timeout: 5_000 }, ); diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts index 994d211cc18..d1b64f37156 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts @@ -21,9 +21,12 @@ test("user can log in", async ({ page }) => { // Filter by the popupRender wrapper class to disambiguate from other // ant-dropdown popups. - const popup = page.locator(".ant-dropdown:visible").filter({ - has: page.locator(".bg-white.rounded-lg.shadow-lg"), - }).first(); + const popup = page + .locator(".ant-dropdown:visible") + .filter({ + has: page.locator(".bg-white.rounded-lg.shadow-lg"), + }) + .first(); await expect(popup).toBeVisible({ timeout: 5_000 }); await expect(popup.getByText("Admin", { exact: true })).toBeVisible({ timeout: 5_000 }); await expect(popup.getByText("default_user_id", { exact: true })).toBeVisible({ timeout: 5_000 }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts index f953a82daaa..22ba85956da 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts @@ -54,9 +54,7 @@ test.describe("MCP Servers", () => { // the MCP servers table so the form modal's `server_name` input — which // still holds the timestamped value during its close animation — can't // satisfy the assertion before the server actually lands in the list. - await expect(page.getByText("MCP Server created successfully").first()) - .toBeVisible({ timeout: 15_000 }); - await expect(page.locator("table tbody").getByText(uniqueName).first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText("MCP Server created successfully").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.locator("table tbody").getByText(uniqueName).first()).toBeVisible({ timeout: 10_000 }); }); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts index ada4dfb735e..ca9c35ce722 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts @@ -31,8 +31,9 @@ test.describe("AI Hub (internal admin view)", () => { // Submit await modal.getByRole("button", { name: "Make Public" }).click(); - await expect(page.getByText(/Successfully made .* model group\(s\) public/i).first()) - .toBeVisible({ timeout: 15_000 }); + await expect(page.getByText(/Successfully made .* model group\(s\) public/i).first()).toBeVisible({ + timeout: 15_000, + }); }); test("AI Hub tab list renders Model Hub, Agent Hub, MCP Hub and Skill Hub", async ({ page }) => { diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts index bb53fb7a23b..17ff1fc3f83 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -154,10 +154,7 @@ test.describe("Add Model", () => { // The Team-BYOK switch is gated on `premiumUser` — without a license set // for the proxy under test, the toggle is disabled and this manual-QA // step cannot be exercised. - test.skip( - !process.env.LITELLM_LICENSE, - "LITELLM_LICENSE not set in test env — Team-BYOK switch is disabled", - ); + test.skip(!process.env.LITELLM_LICENSE, "LITELLM_LICENSE not set in test env — Team-BYOK switch is disabled"); // Make the test idempotent across retries and local reruns: delete any // Cohere model already scoped to the e2e team before we start, and again @@ -170,10 +167,11 @@ test.describe("Add Model", () => { const res = await request.get("/v2/model/info", { headers: auth }); if (!res.ok()) return; const body = await res.json(); - const matches: Array<{ id: string }> = (body?.data ?? []).filter((m: any) => - typeof m?.model_name === "string" && - m.model_name.startsWith("cohere") && - m?.model_info?.team_id === E2E_TEAM_CRUD_ID, + const matches: Array<{ id: string }> = (body?.data ?? []).filter( + (m: any) => + typeof m?.model_name === "string" && + m.model_name.startsWith("cohere") && + m?.model_info?.team_id === E2E_TEAM_CRUD_ID, ); for (const m of matches) { await request.post("/model/delete", { headers: auth, data: { id: m.id } }); @@ -208,9 +206,7 @@ test.describe("Add Model", () => { const teamDropdown = page.getByTestId("team-dropdown"); await expect(teamDropdown).toBeVisible({ timeout: 5_000 }); await teamDropdown.click(); - const teamOption = page.locator(".ant-select-dropdown:visible") - .getByText(E2E_TEAM_CRUD_ID) - .first(); + const teamOption = page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ID).first(); await expect(teamOption).toBeVisible({ timeout: 5_000 }); await teamOption.click(); @@ -219,8 +215,9 @@ test.describe("Add Model", () => { // Scope the success toast to antd's notification container so a stale // success message from an earlier test in the same context can't satisfy // the assertion. - await expect(page.locator(".ant-notification").getByText("created successfully").last()) - .toBeVisible({ timeout: 15_000 }); + await expect(page.locator(".ant-notification").getByText("created successfully").last()).toBeVisible({ + timeout: 15_000, + }); // Verify the model is now in All Models with the team_id attached. The // Models table renders team-scoped models with the team id in the row. @@ -237,16 +234,16 @@ test.describe("Add Model", () => { // Confirm the search returned at least one result — gives a clear // failure message when the table is empty instead of timing out on a // row assertion. - await expect(page.getByTestId("models-results-count")).toHaveText( - /Showing \d+ - \d+ of \d+ results/, - { timeout: 15_000 }, - ); + await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + timeout: 15_000, + }); // Stronger than "alias appears somewhere in tbody" — pin the assertion // to a single row that has BOTH the cohere model_name AND the seeded // team alias, so a stale cohere row from "Add wildcard route" (no team) // can't satisfy the check. - const teamCohereRow = page.locator("table tbody tr") + const teamCohereRow = page + .locator("table tbody tr") .filter({ hasText: "cohere/" }) .filter({ hasText: E2E_TEAM_CRUD_ALIAS }); await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts index d21192d237d..877c7f8c555 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts @@ -62,9 +62,7 @@ test.describe("Clear custom pricing on a deployment", () => { } }); - test("UI sends null for cleared pricing and backend removes the override", async ({ - page, - }) => { + test("UI sends null for cleared pricing and backend removes the override", async ({ page }) => { // Navigate to the model detail view. await page.goto("/ui"); await page.getByText("Models + Endpoints").click(); @@ -97,34 +95,24 @@ test.describe("Clear custom pricing on a deployment", () => { // Capture the outgoing PATCH so we can assert the UI sends explicit nulls. const patchPromise = page.waitForRequest( - (req) => - req.method() === "PATCH" && - req.url().includes(`/model/${createdModelId}/update`) + (req) => req.method() === "PATCH" && req.url().includes(`/model/${createdModelId}/update`), ); await page.getByRole("button", { name: "Save Changes" }).click(); const patchReq = await patchPromise; const patchBody = JSON.parse(patchReq.postData() ?? "{}"); - expect( - patchBody.litellm_params.input_cost_per_token, - "UI sends explicit null for cleared input cost" - ).toBeNull(); - expect( - patchBody.litellm_params.output_cost_per_token, - "UI sends explicit null for cleared output cost" - ).toBeNull(); + expect(patchBody.litellm_params.input_cost_per_token, "UI sends explicit null for cleared input cost").toBeNull(); + expect(patchBody.litellm_params.output_cost_per_token, "UI sends explicit null for cleared output cost").toBeNull(); expect( patchBody.litellm_params.cache_read_input_token_cost, - "UI sends explicit null for cleared cache_read cost" + "UI sends explicit null for cleared cache_read cost", ).toBeNull(); expect( patchBody.litellm_params.cache_creation_input_token_cost, - "UI sends explicit null for cleared cache_write cost" + "UI sends explicit null for cleared cache_write cost", ).toBeNull(); // Success toast confirms the save was accepted. - await expect( - page.getByText("Model settings updated successfully") - ).toBeVisible({ timeout: 10_000 }); + await expect(page.getByText("Model settings updated successfully")).toBeVisible({ timeout: 10_000 }); // Verify via the management API: the user-set rate is gone from both blobs. // The cost-map may synthesize a default for known providers in the response, @@ -132,46 +120,40 @@ test.describe("Clear custom pricing on a deployment", () => { // undefined. const infoRes = await page.request.get( `/v2/model/info?include_team_models=true&page=1&size=100&modelId=${createdModelId}`, - { headers: { Authorization: `Bearer ${masterKey}` } } + { headers: { Authorization: `Bearer ${masterKey}` } }, ); expect(infoRes.ok()).toBe(true); const infoBody = await infoRes.json(); - const row = (infoBody.data ?? infoBody).find?.( - (m: any) => m?.model_info?.id === createdModelId - ); + const row = (infoBody.data ?? infoBody).find?.((m: any) => m?.model_info?.id === createdModelId); expect(row, "model info row").toBeTruthy(); - expect( - "input_cost_per_token" in row.litellm_params, - "litellm_params.input_cost_per_token key removed" - ).toBe(false); - expect( - "output_cost_per_token" in row.litellm_params, - "litellm_params.output_cost_per_token key removed" - ).toBe(false); + expect("input_cost_per_token" in row.litellm_params, "litellm_params.input_cost_per_token key removed").toBe(false); + expect("output_cost_per_token" in row.litellm_params, "litellm_params.output_cost_per_token key removed").toBe( + false, + ); expect( "cache_read_input_token_cost" in row.litellm_params, - "litellm_params.cache_read_input_token_cost key removed" + "litellm_params.cache_read_input_token_cost key removed", ).toBe(false); expect( "cache_creation_input_token_cost" in row.litellm_params, - "litellm_params.cache_creation_input_token_cost key removed" + "litellm_params.cache_creation_input_token_cost key removed", ).toBe(false); expect( row.model_info.input_cost_per_token, - "model_info.input_cost_per_token no longer the seeded override" + "model_info.input_cost_per_token no longer the seeded override", ).not.toBe(SEED_INPUT_PER_TOKEN); expect( row.model_info.output_cost_per_token, - "model_info.output_cost_per_token no longer the seeded override" + "model_info.output_cost_per_token no longer the seeded override", ).not.toBe(SEED_OUTPUT_PER_TOKEN); expect( row.model_info.cache_read_input_token_cost, - "model_info.cache_read_input_token_cost no longer the seeded override" + "model_info.cache_read_input_token_cost no longer the seeded override", ).not.toBe(SEED_CACHE_READ_PER_TOKEN); expect( row.model_info.cache_creation_input_token_cost, - "model_info.cache_creation_input_token_cost no longer the seeded override" + "model_info.cache_creation_input_token_cost no longer the seeded override", ).not.toBe(SEED_CACHE_WRITE_PER_TOKEN); }); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index f56b5875dc6..b8fb95b764d 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -6,15 +6,7 @@ import { menuLabelToPage } from "../../fixtures/menuMappings"; import { navigateToPage } from "../../helpers/navigation"; const sidebarButtons = { - [Role.ProxyAdmin]: [ - "Virtual Keys", - "Playground", - "Models", - "Usage", - "Teams", - "Internal Users", - "AI Hub", - ], + [Role.ProxyAdmin]: ["Virtual Keys", "Playground", "Models", "Usage", "Teams", "Internal Users", "AI Hub"], }; const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }]; diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts index 1e44d9a25a0..644228c5ff9 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts @@ -89,12 +89,8 @@ test.describe("Proxy Admin - Keys", () => { await page.getByRole("spinbutton", { name: "RPM Limit" }).fill("456"); await page.getByRole("button", { name: "Save Changes" }).click(); - await expect( - page.getByRole("paragraph").filter({ hasText: "TPM: 123" }) - ).toBeVisible({ timeout: 10_000 }); - await expect( - page.getByRole("paragraph").filter({ hasText: "RPM: 456" }) - ).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("paragraph").filter({ hasText: "TPM: 123" })).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("paragraph").filter({ hasText: "RPM: 456" })).toBeVisible({ timeout: 10_000 }); }); test("Delete key", async ({ page }) => { diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts index 579b3cede7c..37a0e324f27 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts @@ -14,10 +14,7 @@ import { ADMIN_STORAGE_PATH } from "../../constants"; */ test.describe("Premium license wiring", () => { test("admin session JWT carries premium_user=true when LITELLM_LICENSE is set", () => { - test.skip( - !process.env.LITELLM_LICENSE, - "LITELLM_LICENSE not set in test env — proxy is running unlicensed", - ); + test.skip(!process.env.LITELLM_LICENSE, "LITELLM_LICENSE not set in test env — proxy is running unlicensed"); const storage = JSON.parse(fs.readFileSync(ADMIN_STORAGE_PATH, "utf-8")); const tokenCookie = storage.cookies?.find((c: { name: string }) => c.name === "token"); @@ -28,9 +25,7 @@ test.describe("Premium license wiring", () => { const jwtParts = tokenCookie.value.split("."); expect(jwtParts.length, "token cookie is not a 3-part JWT").toBe(3); const [, payloadB64] = jwtParts; - const payload = JSON.parse( - Buffer.from(payloadB64, "base64url").toString("utf-8"), - ); + const payload = JSON.parse(Buffer.from(payloadB64, "base64url").toString("utf-8")); expect(payload.premium_user).toBe(true); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts index 4774b50dbc5..b30bb8aca7b 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts @@ -19,7 +19,10 @@ test.describe("Proxy Admin - Teams", () => { const uniqueAlias = `e2e-created-team-${Date.now()}`; // Click the Create Team button — accessible name includes "Create Team" - await page.getByRole("button", { name: /Create Team/i }).first().click(); + await page + .getByRole("button", { name: /Create Team/i }) + .first() + .click(); // Wait for the Create Team modal const dialog = page.locator(".ant-modal:visible"); @@ -157,8 +160,9 @@ test.describe("Proxy Admin - Teams", () => { await page.getByRole("button", { name: "Save Changes" }).click(); - await expect(page.getByText(/Team settings updated|updated successfully/i).first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText(/Team settings updated|updated successfully/i).first()).toBeVisible({ + timeout: 10_000, + }); } finally { // Leave the team in its seeded state for any subsequent test or rerun. await restore(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts index 8dd5571f7af..98b86ec9b11 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts @@ -86,17 +86,16 @@ test.describe("Router Settings - Fallbacks", () => { await modal.getByRole("button", { name: /Save All Configurations/i }).click(); // Success toast - await expect(page.getByText(/fallback configuration\(s\) added successfully/i).first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText(/fallback configuration\(s\) added successfully/i).first()).toBeVisible({ + timeout: 10_000, + }); // Modal closes, and a single row contains BOTH the primary and the fallback // model — stronger than asserting each name appears somewhere in tbody, // which could be satisfied by leftover rows from prior runs. await expect(modal).not.toBeVisible({ timeout: 5_000 }); - const newRow = page.locator("table tbody tr") - .filter({ hasText: PRIMARY }) - .filter({ hasText: FALLBACK }); + const newRow = page.locator("table tbody tr").filter({ hasText: PRIMARY }).filter({ hasText: FALLBACK }); await expect(newRow).toHaveCount(1, { timeout: 10_000 }); }); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts index 1612e6929bd..18b43ec89b2 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts @@ -28,13 +28,11 @@ test.describe("Team Admin", () => { await clickTeamId(page, E2E_TEAM_CRUD_ID); await page.getByRole("tab", { name: "Virtual Keys" }).click(); - await expect(page.getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()).toBeVisible({ timeout: 10_000 }); // And from the global Virtual Keys page, the same key should be visible. await navigateToPage(page, Page.ApiKeys); - await expect(page.getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()).toBeVisible({ timeout: 10_000 }); }); test("Team admin can add a member to their team", async ({ page }) => { @@ -60,8 +58,7 @@ test.describe("Team Admin", () => { await modal.getByRole("button", { name: /Add Member/i }).click(); - await expect(page.getByText("Team member added successfully").first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText("Team member added successfully").first()).toBeVisible({ timeout: 10_000 }); }); test("Team admin can remove a member from their team", async ({ page }) => { @@ -82,8 +79,7 @@ test.describe("Team Admin", () => { await expect(modal).toBeVisible({ timeout: 5_000 }); await modal.getByRole("button", { name: /^Delete$/ }).click(); - await expect(page.getByText("Team member removed successfully").first()) - .toBeVisible({ timeout: 10_000 }); + await expect(page.getByText("Team member removed successfully").first()).toBeVisible({ timeout: 10_000 }); }); test("Team admin can create a team key with All Team Models", async ({ page }) => { diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index e93d1997d62..9971ed779ee 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,18 +1,9 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", "entry": ["scripts/**/*.ts"], - "project": [ - "src/**/*.{ts,tsx}", - "tests/**/*.{ts,tsx}", - "scripts/**/*.ts", - "e2e_tests/**/*.ts" - ], + "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.ts", "e2e_tests/**/*.ts"], "playwright": { "config": "e2e_tests/playwright.config.ts", - "entry": [ - "e2e_tests/**/*.spec.ts", - "e2e_tests/**/*.setup.ts", - "e2e_tests/globalSetup.ts" - ] + "entry": ["e2e_tests/**/*.spec.ts", "e2e_tests/**/*.setup.ts", "e2e_tests/globalSetup.ts"] } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/README.md b/ui/litellm-dashboard/src/app/(dashboard)/README.md index c913431fc5b..920ea5b4258 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/README.md +++ b/ui/litellm-dashboard/src/app/(dashboard)/README.md @@ -2,7 +2,7 @@ The LiteLLM UI is currently being refactored/rewritten to reduce development friction. Please read this document to understand what's expected for new contributions. -The project follows strict NextJS file structure. All pages on the site (determined by the sidebar) are contained in their own folder, and routing is automatically handled by NextJS based on the file structure. +The project follows strict NextJS file structure. All pages on the site (determined by the sidebar) are contained in their own folder, and routing is automatically handled by NextJS based on the file structure. For example, NextJS will automatically render the admin settings page when the user visits `/settings/admin-settings` @@ -16,7 +16,9 @@ For example, NextJS will automatically render the admin settings page when the u You can use parenthesis around directory names to hide them from the user route, for example `(dashboard)`, while still getting the benefits of `layout` and file structure. ### File Structure + Every page must follow the following file structure pattern. + ``` ├── teams │   ├── TeamsView.tsx @@ -34,11 +36,11 @@ Every page must follow the following file structure pattern. │   └── page.tsx ``` -### Component Files +### Component Files All component files should ideally be as dumb as possible. Their only job should be to take the data they need from hooks or props and render them to the UI. If a component file becomes too large (over `300` lines or so), **please break it down** into smaller components. -A component should only be placed where it will be used. For example, if a component will only be used by the `teams` page, it should belong in the `teams/components` folder. +A component should only be placed where it will be used. For example, if a component will only be used by the `teams` page, it should belong in the `teams/components` folder. **Common components should be moved to the lowest common ancestor components folder.** diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx index 27a6e6c13be..90f498912a8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx @@ -464,7 +464,6 @@ const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelect /> {isAdminRole(userRole) && !collapsed && } - ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx index 49e6569f1a7..1e091314ecd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx @@ -31,13 +31,15 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side console.log("[SidebarProvider] Fetching UI settings from /get/ui_settings"); const settings = await getUISettings(accessToken); console.log("[SidebarProvider] UI settings response:", settings); - + // API returns 'values' not 'settings' if (settings?.values?.enabled_ui_pages_internal_users !== undefined) { console.log("[SidebarProvider] Setting enabled pages:", settings.values.enabled_ui_pages_internal_users); setEnabledPagesInternalUsers(settings.values.enabled_ui_pages_internal_users); } else { - console.log("[SidebarProvider] No enabled_ui_pages_internal_users in response (all pages visible by default)"); + console.log( + "[SidebarProvider] No enabled_ui_pages_internal_users in response (all pages visible by default)", + ); } if (settings?.values?.enable_projects_ui !== undefined) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails.ts index c0379b25321..3dcf73388a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails.ts @@ -1,20 +1,12 @@ import { useQuery, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import { all_admin_roles } from "@/utils/roles"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { AccessGroupResponse, accessGroupKeys } from "./useAccessGroups"; // ── Fetch function ─────────────────────────────────────────────────────────── -const fetchAccessGroupDetails = async ( - accessToken: string, - accessGroupId: string, -): Promise => { +const fetchAccessGroupDetails = async (accessToken: string, accessGroupId: string): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/v1/access_group/${encodeURIComponent(accessGroupId)}`; @@ -45,17 +37,13 @@ export const useAccessGroupDetails = (accessGroupId?: string) => { return useQuery({ queryKey: accessGroupKeys.detail(accessGroupId!), queryFn: async () => fetchAccessGroupDetails(accessToken!, accessGroupId!), - enabled: - Boolean(accessToken && accessGroupId) && - all_admin_roles.includes(userRole || ""), + enabled: Boolean(accessToken && accessGroupId) && all_admin_roles.includes(userRole || ""), // Seed from the list cache when available initialData: () => { if (!accessGroupId) return undefined; - const groups = queryClient.getQueryData( - accessGroupKeys.list({}), - ); + const groups = queryClient.getQueryData(accessGroupKeys.list({})); return groups?.find((g) => g.access_group_id === accessGroupId); }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroups.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroups.ts index 215b555fcf9..9f306c21459 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroups.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useAccessGroups.ts @@ -1,11 +1,6 @@ import { useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import { all_admin_roles } from "@/utils/roles"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -32,9 +27,7 @@ export const accessGroupKeys = createQueryKeys("accessGroups"); // ── Fetch function ─────────────────────────────────────────────────────────── -const fetchAccessGroups = async ( - accessToken: string, -): Promise => { +const fetchAccessGroups = async (accessToken: string): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/v1/access_group`; @@ -64,7 +57,6 @@ export const useAccessGroups = () => { return useQuery({ queryKey: accessGroupKeys.list({}), queryFn: async () => fetchAccessGroups(accessToken!), - enabled: - Boolean(accessToken) && all_admin_roles.includes(userRole || ""), + enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""), }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useCreateAccessGroup.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useCreateAccessGroup.ts index 7ea5a813462..5efa2da6557 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useCreateAccessGroup.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useCreateAccessGroup.ts @@ -1,10 +1,5 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { AccessGroupResponse, accessGroupKeys } from "./useAccessGroups"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup.ts index 5df5960ce0a..01e317f6613 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup.ts @@ -1,19 +1,11 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { accessGroupKeys } from "./useAccessGroups"; // ── Fetch function ─────────────────────────────────────────────────────────── -const deleteAccessGroup = async ( - accessToken: string, - accessGroupId: string, -): Promise => { +const deleteAccessGroup = async (accessToken: string, accessGroupId: string): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/v1/access_group/${encodeURIComponent(accessGroupId)}`; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useEditAccessGroup.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useEditAccessGroup.ts index 5dc2252f640..7dd85ae93dc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useEditAccessGroup.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/accessGroups/useEditAccessGroup.ts @@ -1,10 +1,5 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { AccessGroupResponse, accessGroupKeys } from "./useAccessGroups"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.test.ts index 8334aea56e7..f370e4d6d6e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.test.ts @@ -4,27 +4,22 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import React, { ReactNode } from "react"; import { useCloudZeroCreate } from "./useCloudZeroCreate"; -const { - mockProxyBaseUrl, - mockAccessToken, - mockHeaderName, - mockGetProxyBaseUrl, - mockGetGlobalLitellmHeaderName, -} = vi.hoisted(() => { - const mockProxyBaseUrl = "https://proxy.example.com"; - const mockAccessToken = "test-access-token"; - const mockHeaderName = "X-LiteLLM-API-Key"; - const mockGetProxyBaseUrl = vi.fn(() => mockProxyBaseUrl); - const mockGetGlobalLitellmHeaderName = vi.fn(() => mockHeaderName); +const { mockProxyBaseUrl, mockAccessToken, mockHeaderName, mockGetProxyBaseUrl, mockGetGlobalLitellmHeaderName } = + vi.hoisted(() => { + const mockProxyBaseUrl = "https://proxy.example.com"; + const mockAccessToken = "test-access-token"; + const mockHeaderName = "X-LiteLLM-API-Key"; + const mockGetProxyBaseUrl = vi.fn(() => mockProxyBaseUrl); + const mockGetGlobalLitellmHeaderName = vi.fn(() => mockHeaderName); - return { - mockProxyBaseUrl, - mockAccessToken, - mockHeaderName, - mockGetProxyBaseUrl, - mockGetGlobalLitellmHeaderName, - }; -}); + return { + mockProxyBaseUrl, + mockAccessToken, + mockHeaderName, + mockGetProxyBaseUrl, + mockGetGlobalLitellmHeaderName, + }; + }); vi.mock("@/components/networking", () => ({ getProxyBaseUrl: mockGetProxyBaseUrl, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.test.ts index 74d657b3e85..b5b903ea620 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.test.ts @@ -4,27 +4,22 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import React, { ReactNode } from "react"; import { useCloudZeroDryRun } from "./useCloudZeroDryRun"; -const { - mockProxyBaseUrl, - mockAccessToken, - mockHeaderName, - mockGetProxyBaseUrl, - mockGetGlobalLitellmHeaderName, -} = vi.hoisted(() => { - const mockProxyBaseUrl = "https://proxy.example.com"; - const mockAccessToken = "test-access-token"; - const mockHeaderName = "X-LiteLLM-API-Key"; - const mockGetProxyBaseUrl = vi.fn(() => mockProxyBaseUrl); - const mockGetGlobalLitellmHeaderName = vi.fn(() => mockHeaderName); +const { mockProxyBaseUrl, mockAccessToken, mockHeaderName, mockGetProxyBaseUrl, mockGetGlobalLitellmHeaderName } = + vi.hoisted(() => { + const mockProxyBaseUrl = "https://proxy.example.com"; + const mockAccessToken = "test-access-token"; + const mockHeaderName = "X-LiteLLM-API-Key"; + const mockGetProxyBaseUrl = vi.fn(() => mockProxyBaseUrl); + const mockGetGlobalLitellmHeaderName = vi.fn(() => mockHeaderName); - return { - mockProxyBaseUrl, - mockAccessToken, - mockHeaderName, - mockGetProxyBaseUrl, - mockGetGlobalLitellmHeaderName, - }; -}); + return { + mockProxyBaseUrl, + mockAccessToken, + mockHeaderName, + mockGetProxyBaseUrl, + mockGetGlobalLitellmHeaderName, + }; + }); vi.mock("@/components/networking", () => ({ getProxyBaseUrl: mockGetProxyBaseUrl, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.test.ts index 72a1cfd24aa..3c44d75dd06 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.test.ts @@ -4,27 +4,22 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import React, { ReactNode } from "react"; import { useCloudZeroExport } from "./useCloudZeroExport"; -const { - mockProxyBaseUrl, - mockAccessToken, - mockHeaderName, - mockGetProxyBaseUrl, - mockGetGlobalLitellmHeaderName, -} = vi.hoisted(() => { - const mockProxyBaseUrl = "https://proxy.example.com"; - const mockAccessToken = "test-access-token"; - const mockHeaderName = "X-LiteLLM-API-Key"; - const mockGetProxyBaseUrl = vi.fn(() => mockProxyBaseUrl); - const mockGetGlobalLitellmHeaderName = vi.fn(() => mockHeaderName); +const { mockProxyBaseUrl, mockAccessToken, mockHeaderName, mockGetProxyBaseUrl, mockGetGlobalLitellmHeaderName } = + vi.hoisted(() => { + const mockProxyBaseUrl = "https://proxy.example.com"; + const mockAccessToken = "test-access-token"; + const mockHeaderName = "X-LiteLLM-API-Key"; + const mockGetProxyBaseUrl = vi.fn(() => mockProxyBaseUrl); + const mockGetGlobalLitellmHeaderName = vi.fn(() => mockHeaderName); - return { - mockProxyBaseUrl, - mockAccessToken, - mockHeaderName, - mockGetProxyBaseUrl, - mockGetGlobalLitellmHeaderName, - }; -}); + return { + mockProxyBaseUrl, + mockAccessToken, + mockHeaderName, + mockGetProxyBaseUrl, + mockGetGlobalLitellmHeaderName, + }; + }); vi.mock("@/components/networking", () => ({ getProxyBaseUrl: mockGetProxyBaseUrl, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts index 39afd044097..2c1fd29f61c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts @@ -13,11 +13,7 @@ describe("createQueryKeys", () => { }); it("should generate a list key with params", () => { - expect(keys.list({ page: 1, limit: 10 })).toEqual([ - "books", - "list", - { params: { page: 1, limit: 10 } }, - ]); + expect(keys.list({ page: 1, limit: 10 })).toEqual(["books", "list", { params: { page: 1, limit: 10 } }]); }); it("should generate a list key with undefined params when none provided", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/hashicorpVaultApi.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/hashicorpVaultApi.ts index edf18860ec1..2af0f118500 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/hashicorpVaultApi.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/hashicorpVaultApi.ts @@ -2,9 +2,7 @@ import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage } from export const getHashicorpVaultConfig = async (accessToken: string) => { const proxyBaseUrl = getProxyBaseUrl(); - const url = proxyBaseUrl - ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` - : `/config_overrides/hashicorp_vault`; + const url = proxyBaseUrl ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` : `/config_overrides/hashicorp_vault`; const response = await fetch(url, { method: "GET", headers: { @@ -20,14 +18,9 @@ export const getHashicorpVaultConfig = async (accessToken: string) => { return data; }; -export const updateHashicorpVaultConfig = async ( - accessToken: string, - config: Record, -) => { +export const updateHashicorpVaultConfig = async (accessToken: string, config: Record) => { const proxyBaseUrl = getProxyBaseUrl(); - const url = proxyBaseUrl - ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` - : `/config_overrides/hashicorp_vault`; + const url = proxyBaseUrl ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` : `/config_overrides/hashicorp_vault`; const response = await fetch(url, { method: "POST", headers: { @@ -47,9 +40,7 @@ export const updateHashicorpVaultConfig = async ( export const deleteHashicorpVaultConfig = async (accessToken: string) => { const proxyBaseUrl = getProxyBaseUrl(); - const url = proxyBaseUrl - ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` - : `/config_overrides/hashicorp_vault`; + const url = proxyBaseUrl ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` : `/config_overrides/hashicorp_vault`; const response = await fetch(url, { method: "DELETE", headers: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts index b1896eda0e6..8db520eecd0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts @@ -289,11 +289,7 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data?.globalGuardrailNames).toEqual( - new Set(["global-guard-a", "global-guard-b"]), - ); - expect(result.current.data?.optionalGuardrailNames).toEqual( - new Set(["optional-guard-a", "optional-guard-b"]), - ); + expect(result.current.data?.globalGuardrailNames).toEqual(new Set(["global-guard-a", "global-guard-b"])); + expect(result.current.data?.optionalGuardrailNames).toEqual(new Set(["optional-guard-a", "optional-guard-b"])); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts index 3135e8326fc..edbcbdfe170 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts @@ -1,10 +1,5 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { createQueryKeys } from "../common/queryKeysFactory"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts index 5838dbd0ee6..3b79e5c7643 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts @@ -1,8 +1,5 @@ import { useQuery, UseQueryResult } from "@tanstack/react-query"; -import { - getGlobalLitellmHeaderName, - getProxyBaseUrl, -} from "@/components/networking"; +import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; import { createQueryKeys } from "../common/queryKeysFactory"; const healthReadinessDetailsKeys = createQueryKeys("healthReadinessDetails"); @@ -18,9 +15,7 @@ export interface HealthReadinessDetailsResponse { is_detailed_debug?: boolean; } -const fetchHealthReadinessDetails = async ( - accessToken: string, -): Promise => { +const fetchHealthReadinessDetails = async (accessToken: string): Promise => { const baseUrl = getProxyBaseUrl(); const response = await fetch(`${baseUrl}/health/readiness/details`, { method: "GET", @@ -30,9 +25,7 @@ const fetchHealthReadinessDetails = async ( }, }); if (!response.ok) { - throw new Error( - `Failed to fetch health readiness details: ${response.statusText}`, - ); + throw new Error(`Failed to fetch health readiness details: ${response.statusText}`); } return response.json(); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts index 1e1190b12c8..e0140c6a63a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts @@ -128,9 +128,7 @@ describe("useInfiniteKeyAliases", () => { }); it("should fetch the next page when fetchNextPage is called", async () => { - mockKeyAliasesCall - .mockResolvedValueOnce(mockPage1) - .mockResolvedValueOnce(mockPage2); + mockKeyAliasesCall.mockResolvedValueOnce(mockPage1).mockResolvedValueOnce(mockPage2); const wrapper = createWrapper(); const { result } = renderHook(() => useInfiniteKeyAliases(2), { wrapper }); @@ -151,10 +149,10 @@ describe("useInfiniteKeyAliases", () => { it("should include search in query key so search changes refetch from page 1", async () => { const wrapper = createWrapper(); - const { result, rerender } = renderHook( - ({ search }: { search?: string }) => useInfiniteKeyAliases(50, search), - { wrapper, initialProps: { search: undefined } }, - ); + const { result, rerender } = renderHook(({ search }: { search?: string }) => useInfiniteKeyAliases(50, search), { + wrapper, + initialProps: { search: undefined }, + }); await waitFor(() => { expect(result.current.isSuccess).toBe(true); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts index 03e96fe73c4..2b4583ad6b3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts @@ -5,11 +5,7 @@ import useAuthorized from "../useAuthorized"; const infiniteKeyAliasKeys = createQueryKeys("infiniteKeyAliases"); -export const useInfiniteKeyAliases = ( - size: number = 50, - search?: string, - team_id?: string, -) => { +export const useInfiniteKeyAliases = (size: number = 50, search?: string, team_id?: string) => { const { accessToken } = useAuthorized(); return useInfiniteQuery({ queryKey: infiniteKeyAliasKeys.list({ @@ -20,13 +16,7 @@ export const useInfiniteKeyAliases = ( }, }), queryFn: async ({ pageParam }) => { - return await keyAliasesCall( - accessToken!, - pageParam as number, - size, - search, - team_id, - ); + return await keyAliasesCall(accessToken!, pageParam as number, size, search, team_id); }, initialPageParam: 1, getNextPageParam: (lastPage) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts index 80cb69495da..1e700e572d0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -410,10 +410,7 @@ describe("useKeys", () => { }), }); - const { result } = renderHook( - () => useKeys(1, 10, { projectID: "project-1" }), - { wrapper }, - ); + const { result } = renderHook(() => useKeys(1, 10, { projectID: "project-1" }), { wrapper }); await waitFor(() => { expect(result.current.isLoading).toBe(false); @@ -436,10 +433,7 @@ describe("useKeys", () => { }), }); - const { result } = renderHook( - () => useKeys(1, 10, { projectID: "project-1", teamID: "team-1" }), - { wrapper }, - ); + const { result } = renderHook(() => useKeys(1, 10, { projectID: "project-1", teamID: "team-1" }), { wrapper }); await waitFor(() => { expect(result.current.isLoading).toBe(false); @@ -456,10 +450,7 @@ describe("useKeys", () => { json: async () => mockKeysResponse, }); - const { result } = renderHook( - () => useKeys(1, 10, { projectID: null }), - { wrapper }, - ); + const { result } = renderHook(() => useKeys(1, 10, { projectID: null }), { wrapper }); await waitFor(() => { expect(result.current.isLoading).toBe(false); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index fbe5eccb75a..4a04c541d1a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -1,11 +1,6 @@ import { keepPreviousData, useQuery, UseQueryResult } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -43,18 +38,13 @@ export interface KeyListCallOptions { status?: string | null; } -const keyListCall = async ( - accessToken: string, - page: number, - pageSize: number, - options: KeyListCallOptions = {}, -) => { +const keyListCall = async (accessToken: string, page: number, pageSize: number, options: KeyListCallOptions = {}) => { /** * Get all available keys on proxy */ try { const baseUrl = getProxyBaseUrl(); - + const params = new URLSearchParams( Object.entries({ team_id: options.teamID, @@ -134,4 +124,4 @@ export const useDeletedKeys = ( staleTime: 30000, // 30 seconds placeholderData: keepPreviousData, }); -}; \ No newline at end of file +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts index a845fc5881a..0265b4dc402 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts @@ -1,10 +1,5 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { keyKeys } from "./useKeys"; @@ -20,10 +15,7 @@ export interface ResetKeySpendResponse { // ── Fetch function ──────────────────────────────────────────────────────────── -export const resetKeySpend = async ( - accessToken: string, - keyToken: string, -): Promise => { +export const resetKeySpend = async (accessToken: string, keyToken: string): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl ? `${baseUrl}/key/${keyToken}/reset_spend` : `/key/${keyToken}/reset_spend`}`; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/logDetails/useLogDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/logDetails/useLogDetails.ts index 6c0f95d5995..5e4757bdb2f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/logDetails/useLogDetails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/logDetails/useLogDetails.ts @@ -10,11 +10,7 @@ import { uiSpendLogDetailsCall } from "@/components/networking"; * @param startTime - The formatted start time for the query * @param enabled - Whether the query should be enabled (e.g., drawer is open) */ -export const useLogDetails = ( - requestId: string | undefined, - startTime: string | undefined, - enabled: boolean, -) => { +export const useLogDetails = (requestId: string | undefined, startTime: string | undefined, enabled: boolean) => { const { accessToken } = useAuthorized(); return useQuery({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts index e91f5aa670b..ad9880d8cac 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts @@ -3,9 +3,7 @@ import { useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import useAuthorized from "../useAuthorized"; -const mcpSemanticFilterSettingsKeys = createQueryKeys( - "mcpSemanticFilterSettings" -); +const mcpSemanticFilterSettingsKeys = createQueryKeys("mcpSemanticFilterSettings"); export const useMCPSemanticFilterSettings = () => { const { accessToken } = useAuthorized(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts index 2062b4f4c29..bc7406599b1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts @@ -2,9 +2,7 @@ import { updateMCPSemanticFilterSettings } from "@/components/networking"; import { useMutation, useQueryClient } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -const mcpSemanticFilterSettingsKeys = createQueryKeys( - "mcpSemanticFilterSettings" -); +const mcpSemanticFilterSettingsKeys = createQueryKeys("mcpSemanticFilterSettings"); export const useUpdateMCPSemanticFilterSettings = (accessToken: string) => { const queryClient = useQueryClient(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.test.ts index 9c555ff1234..65dfd6bf4f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.test.ts @@ -121,4 +121,4 @@ describe("useMCPAccessGroups", () => { expect(result.current.data).toEqual([]); }); -}); \ No newline at end of file +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts index 681bf4161ad..9ad8a6f43fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -24,32 +24,32 @@ export const useMCPServerHealth = () => { refetchInterval: 30000, }); - const recheckServerHealth = useCallback(async (serverId: string) => { - if (!accessToken) return; + const recheckServerHealth = useCallback( + async (serverId: string) => { + if (!accessToken) return; - setRecheckingServerIds((prev) => new Set(prev).add(serverId)); + setRecheckingServerIds((prev) => new Set(prev).add(serverId)); - try { - const result: MCPServerHealth[] = await fetchMCPServerHealth(accessToken, [serverId]); + try { + const result: MCPServerHealth[] = await fetchMCPServerHealth(accessToken, [serverId]); - queryClient.setQueriesData( - { queryKey: mcpServerHealthKeys.lists() }, - (oldData) => { + queryClient.setQueriesData({ queryKey: mcpServerHealthKeys.lists() }, (oldData) => { if (!oldData) return result; return oldData.map((h) => { const updated = result.find((r) => r.server_id === h.server_id); return updated ?? h; }); - }, - ); - } finally { - setRecheckingServerIds((prev) => { - const next = new Set(prev); - next.delete(serverId); - return next; - }); - } - }, [accessToken, queryClient]); + }); + } finally { + setRecheckingServerIds((prev) => { + const next = new Set(prev); + next.delete(serverId); + return next; + }); + } + }, + [accessToken, queryClient], + ); return { ...query, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts index ee03a0ab7c3..52b58f9e318 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts @@ -131,4 +131,4 @@ describe("useMCPServers", () => { expect(result.current.data).toEqual([]); }); -}); \ No newline at end of file +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index fe1afdcc39f..c997f679b2e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -28,7 +28,15 @@ const allProxyModelsKeys = createQueryKeys("allProxyModels"); const selectedTeamModelsKeys = createQueryKeys("selectedTeamModels"); const infiniteModelKeys = createQueryKeys("infiniteModels"); -export const useModelsInfo = (page: number = 1, size: number = 50, search?: string, modelId?: string, teamId?: string, sortBy?: string, sortOrder?: string) => { +export const useModelsInfo = ( + page: number = 1, + size: number = 50, + search?: string, + modelId?: string, + teamId?: string, + sortBy?: string, + sortOrder?: string, +) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ queryKey: modelKeys.list({ @@ -44,7 +52,8 @@ export const useModelsInfo = (page: number = 1, size: number = 50, search?: stri ...(sortOrder && { sortOrder }), }, }), - queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!, page, size, search, modelId, teamId, sortBy, sortOrder), + queryFn: async () => + await modelInfoCall(accessToken!, userId!, userRole!, page, size, search, modelId, teamId, sortBy, sortOrder), enabled: Boolean(accessToken && userId && userRole), }); }; @@ -76,10 +85,7 @@ export const useSelectedTeamModels = (teamID: string | null) => { }); }; -export const useInfiniteModelInfo = ( - size: number = 50, - search?: string, -) => { +export const useInfiniteModelInfo = (size: number = 50, search?: string) => { const { accessToken, userId, userRole } = useAuthorized(); return useInfiniteQuery({ queryKey: infiniteModelKeys.list({ @@ -91,14 +97,7 @@ export const useInfiniteModelInfo = ( }, }), queryFn: async ({ pageParam }) => { - return await modelInfoCall( - accessToken!, - userId!, - userRole!, - pageParam as number, - size, - search, - ); + return await modelInfoCall(accessToken!, userId!, userRole!, pageParam as number, size, search); }, initialPageParam: 1, getNextPageParam: (lastPage) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts index 64d950d59ee..110a704725a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts @@ -103,9 +103,7 @@ describe("useCreateProject", () => { const { result } = renderHook(() => useCreateProject(), { wrapper: makeWrapper(queryClient), }); - await expect(result.current.mutateAsync({ team_id: "team-1" })).rejects.toThrow( - "Access token is required" - ); + await expect(result.current.mutateAsync({ team_id: "team-1" })).rejects.toThrow("Access token is required"); expect(global.fetch).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts index e206c770b19..2e67e626936 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts @@ -1,10 +1,5 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { ProjectResponse, projectKeys } from "./useProjects"; @@ -25,10 +20,7 @@ export interface ProjectCreateParams { // ── Fetch function ─────────────────────────────────────────────────────────── -const createProject = async ( - accessToken: string, - params: ProjectCreateParams, -): Promise => { +const createProject = async (accessToken: string, params: ProjectCreateParams): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/project/new`; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts index 85a9f3e0b10..beaad13ce2a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts @@ -80,9 +80,7 @@ describe("useDeleteProject", () => { const { result } = renderHook(() => useDeleteProject(), { wrapper: makeWrapper(queryClient), }); - await expect(result.current.mutateAsync(["proj-1"])).rejects.toThrow( - "Access token is required" - ); + await expect(result.current.mutateAsync(["proj-1"])).rejects.toThrow("Access token is required"); expect(global.fetch).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.ts index 5abf9e03be2..04f2c547eef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.ts @@ -1,19 +1,11 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { projectKeys } from "./useProjects"; // ── Fetch function ─────────────────────────────────────────────────────────── -const deleteProjects = async ( - accessToken: string, - projectIds: string[], -): Promise => { +const deleteProjects = async (accessToken: string, projectIds: string[]): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/project/delete`; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts index 1d35ac1bf70..037baa18692 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts @@ -1,20 +1,12 @@ import { useQuery, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import { all_admin_roles } from "@/utils/roles"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { ProjectResponse, projectKeys } from "./useProjects"; // ── Fetch function ─────────────────────────────────────────────────────────── -const fetchProjectDetails = async ( - accessToken: string, - projectId: string, -): Promise => { +const fetchProjectDetails = async (accessToken: string, projectId: string): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/project/info?project_id=${encodeURIComponent(projectId)}`; @@ -45,17 +37,13 @@ export const useProjectDetails = (projectId?: string) => { return useQuery({ queryKey: projectKeys.detail(projectId!), queryFn: async () => fetchProjectDetails(accessToken!, projectId!), - enabled: - Boolean(accessToken && projectId) && - all_admin_roles.includes(userRole || ""), + enabled: Boolean(accessToken && projectId) && all_admin_roles.includes(userRole || ""), // Seed from the list cache when available initialData: () => { if (!projectId) return undefined; - const projects = queryClient.getQueryData( - projectKeys.list({}), - ); + const projects = queryClient.getQueryData(projectKeys.list({})); return projects?.find((p) => p.project_id === projectId); }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts index 79976f54626..7bdc8a4fe6d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts @@ -1,11 +1,6 @@ import { useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { all_admin_roles } from "@/utils/roles"; @@ -49,9 +44,7 @@ export const projectKeys = createQueryKeys("projects"); // ── Fetch function ─────────────────────────────────────────────────────────── -const fetchProjects = async ( - accessToken: string, -): Promise => { +const fetchProjects = async (accessToken: string): Promise => { const baseUrl = getProxyBaseUrl(); const url = `${baseUrl}/project/list`; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts index 31d1a5fb352..9e752ac098a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts @@ -108,9 +108,9 @@ describe("useUpdateProject", () => { const { result } = renderHook(() => useUpdateProject(), { wrapper: makeWrapper(queryClient), }); - await expect( - result.current.mutateAsync({ projectId: "proj-1", params: {} }) - ).rejects.toThrow("Access token is required"); + await expect(result.current.mutateAsync({ projectId: "proj-1", params: {} })).rejects.toThrow( + "Access token is required", + ); expect(global.fetch).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts index 2042c8fc7cd..6d8c2d9d4f8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts @@ -1,10 +1,5 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { ProjectResponse, projectKeys } from "./useProjects"; @@ -58,11 +53,7 @@ export const useUpdateProject = () => { const { accessToken } = useAuthorized(); const queryClient = useQueryClient(); - return useMutation< - ProjectResponse, - Error, - { projectId: string; params: ProjectUpdateParams } - >({ + return useMutation({ mutationFn: async ({ projectId, params }) => { if (!accessToken) { throw new Error("Access token is required"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts index 6ff784ebd90..dd69e8c8791 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts @@ -54,7 +54,7 @@ describe("useStoreModelInDB", () => { field_value: true, config_type: "general_settings", }), - }) + }), ); }); @@ -80,15 +80,12 @@ describe("useStoreModelInDB", () => { field_value: false, config_type: "general_settings", }), - }) + }), ); }); it("should throw error when access token is missing", async () => { - vi.spyOn( - await import("../useAuthorized"), - "default" - ).mockReturnValue({ + vi.spyOn(await import("../useAuthorized"), "default").mockReturnValue({ accessToken: null, userRole: null, userId: null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts index e6efbd724cd..27e375c265d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts @@ -12,7 +12,7 @@ export interface StoreModelInDBResponse { const performStoreModelInDB = async ( accessToken: string, - params: StoreModelInDBParams + params: StoreModelInDBParams, ): Promise => { const proxyBaseUrl = getProxyBaseUrl(); const url = proxyBaseUrl ? `${proxyBaseUrl}/config/field/update` : `/config/field/update`; @@ -41,11 +41,7 @@ const performStoreModelInDB = async ( return data; }; -export const useStoreModelInDB = (): UseMutationResult< - StoreModelInDBResponse, - Error, - StoreModelInDBParams -> => { +export const useStoreModelInDB = (): UseMutationResult => { const { accessToken } = useAuthorized(); return useMutation({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts index 67b52997a01..88a37b30291 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts @@ -14,7 +14,7 @@ export interface StoreRequestInSpendLogsResponse { const performStoreRequestInSpendLogs = async ( accessToken: string, - params: StoreRequestInSpendLogsParams + params: StoreRequestInSpendLogsParams, ): Promise => { const proxyBaseUrl = getProxyBaseUrl(); const url = proxyBaseUrl ? `${proxyBaseUrl}/config/update` : `/config/update`; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts index 217ca426c25..20f034ada36 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -423,7 +423,7 @@ describe("useTeam", () => { // This tests the defensive error path in queryFn (lines 111-112) // The enabled check prevents queryFn from running, but we can test the defensive code // by manually constructing and calling the queryFn logic - + // Set up mocks mockUseAuthorized.mockReturnValue({ accessToken: null, // Missing accessToken @@ -438,24 +438,24 @@ describe("useTeam", () => { // Import useQueryClient to get access to query client const { useQueryClient } = await import("@tanstack/react-query"); - + // Manually test the queryFn logic by calling it directly // This simulates what would happen if enabled check was bypassed const testQueryFn = async () => { const { accessToken } = mockUseAuthorized(); const teamId = "team-1"; - + // This is the defensive check from lines 111-112 if (!accessToken || !teamId) { throw new Error("Missing auth or teamId"); } - + return teamInfoCall(accessToken, teamId); }; // Test that the error is thrown await expect(testQueryFn()).rejects.toThrow("Missing auth or teamId"); - + // Also test with missing teamId mockUseAuthorized.mockReturnValue({ accessToken: "test-access-token", @@ -471,11 +471,11 @@ describe("useTeam", () => { const testQueryFnMissingTeamId = async () => { const { accessToken } = mockUseAuthorized(); const teamId = undefined; // Missing teamId - + if (!accessToken || !teamId) { throw new Error("Missing auth or teamId"); } - + return teamInfoCall(accessToken, teamId); }; @@ -736,13 +736,10 @@ describe("useDeletedTeams", () => { json: async () => ({ teams: mockDeletedTeams }), }); - const { result, rerender } = renderHook( - ({ page }) => useDeletedTeams(page, 10, {}), - { - wrapper, - initialProps: { page: 1 }, - }, - ); + const { result, rerender } = renderHook(({ page }) => useDeletedTeams(page, 10, {}), { + wrapper, + initialProps: { page: 1 }, + }); await waitFor(() => { expect(result.current.isSuccess).toBe(true); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index b25b6ce393a..c356434ba04 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -4,12 +4,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchTeams } from "@/app/(dashboard)/networking"; import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory"; import { teamInfoCall } from "@/components/networking"; -import { - getProxyBaseUrl, - getGlobalLitellmHeaderName, - deriveErrorMessage, - handleError, -} from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; export interface TeamsResponse { teams: Team[]; @@ -24,7 +19,6 @@ export interface DeletedTeam extends Team { deleted_by: string; } - export interface TeamListCallOptions { organizationID?: string | null; teamID?: string | null; @@ -47,7 +41,7 @@ export const teamListCall = async ( */ try { const baseUrl = getProxyBaseUrl(); - + const params = new URLSearchParams( Object.entries({ team_id: options.teamID, @@ -128,11 +122,7 @@ export const useTeam = (teamId?: string) => { const infiniteTeamKeys = createQueryKeys("infiniteTeams"); -export const useInfiniteTeams = ( - pageSize: number = 50, - search?: string, - organizationId?: string | null, -) => { +export const useInfiniteTeams = (pageSize: number = 50, search?: string, organizationId?: string | null) => { const { accessToken, userId, userRole } = useAuthorized(); const isAdmin = userRole === "Admin" || userRole === "Admin Viewer"; @@ -174,7 +164,7 @@ const deletedTeamListCall = async ( */ try { const baseUrl = getProxyBaseUrl(); - + const params = new URLSearchParams( Object.entries({ team_id: options.teamID, @@ -211,10 +201,10 @@ const deletedTeamListCall = async ( const data = await response.json(); console.log("/team/list?status=deleted API Response:", data); - + // Extract teams array from response if it's wrapped in a response object // Otherwise return the data directly if it's already an array - if (data && typeof data === 'object' && 'teams' in data) { + if (data && typeof data === "object" && "teams" in data) { return data.teams as DeletedTeam[]; } return data as DeletedTeam[]; @@ -239,4 +229,4 @@ export const useDeletedTeams = ( staleTime: 30000, // 30 seconds placeholderData: keepPreviousData, }); -}; \ No newline at end of file +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index 5178aca0790..94f9d9173f0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -8,7 +8,15 @@ import useAuthorized from "./useAuthorized"; // Unmock useAuthorized to test the actual implementation vi.unmock("@/app/(dashboard)/hooks/useAuthorized"); -const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock, buildLoginUrlWithReturnMock } = vi.hoisted(() => ({ +const { + replaceMock, + clearTokenCookiesMock, + getProxyBaseUrlMock, + getUiConfigMock, + decodeTokenMock, + checkTokenValidityMock, + buildLoginUrlWithReturnMock, +} = vi.hoisted(() => ({ replaceMock: vi.fn(), clearTokenCookiesMock: vi.fn(), getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), @@ -102,7 +110,7 @@ describe("useAuthorized", () => { admin_ui_disabled: false, sso_configured: false, }); - + const decodedPayload = { key: "api-key-123", user_id: "user-1", @@ -112,7 +120,7 @@ describe("useAuthorized", () => { disabled_non_admin_personal_key_creation: false, login_method: "username_password", }; - + decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts index b0a96eff0e7..537e2c5378a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts @@ -36,11 +36,7 @@ const DEFAULT_AUTH = { showSSOBanner: false, }; -const buildUserListResponse = ( - page: number, - totalPages: number, - userCount = 2, -): UserListResponse => ({ +const buildUserListResponse = (page: number, totalPages: number, userCount = 2): UserListResponse => ({ page, page_size: 50, total: totalPages * userCount, @@ -90,13 +86,7 @@ describe("useInfiniteUsers", () => { expect(result.current.data?.pages).toHaveLength(1); expect(result.current.data?.pages[0]).toEqual(mockResponse); - expect(userListCall).toHaveBeenCalledWith( - "test-access-token", - null, - 1, - 50, - null, - ); + expect(userListCall).toHaveBeenCalledWith("test-access-token", null, 1, 50, null); }); it("should use the default page size of 50", async () => { @@ -109,13 +99,7 @@ describe("useInfiniteUsers", () => { expect(result.current.isSuccess).toBe(true); }); - expect(userListCall).toHaveBeenCalledWith( - "test-access-token", - null, - 1, - 50, - null, - ); + expect(userListCall).toHaveBeenCalledWith("test-access-token", null, 1, 50, null); }); it("should use a custom page size when provided", async () => { @@ -131,13 +115,7 @@ describe("useInfiniteUsers", () => { expect(result.current.isSuccess).toBe(true); }); - expect(userListCall).toHaveBeenCalledWith( - "test-access-token", - null, - 1, - customPageSize, - null, - ); + expect(userListCall).toHaveBeenCalledWith("test-access-token", null, 1, customPageSize, null); }); it("should pass searchEmail to userListCall when provided", async () => { @@ -153,13 +131,7 @@ describe("useInfiniteUsers", () => { expect(result.current.isSuccess).toBe(true); }); - expect(userListCall).toHaveBeenCalledWith( - "test-access-token", - null, - 1, - 50, - searchEmail, - ); + expect(userListCall).toHaveBeenCalledWith("test-access-token", null, 1, 50, searchEmail); }); it("should pass null for searchEmail when not provided", async () => { @@ -174,13 +146,7 @@ describe("useInfiniteUsers", () => { expect(result.current.isSuccess).toBe(true); }); - expect(userListCall).toHaveBeenCalledWith( - "test-access-token", - null, - 1, - 50, - null, - ); + expect(userListCall).toHaveBeenCalledWith("test-access-token", null, 1, 50, null); }); it("should fetch the next page when more pages are available", async () => { @@ -209,13 +175,7 @@ describe("useInfiniteUsers", () => { expect(result.current.data?.pages[1]).toEqual(page2); expect(userListCall).toHaveBeenCalledTimes(2); - expect(userListCall).toHaveBeenLastCalledWith( - "test-access-token", - null, - 2, - 50, - null, - ); + expect(userListCall).toHaveBeenLastCalledWith("test-access-token", null, 2, 50, null); }); it("should not have a next page when on the last page", async () => { @@ -275,13 +235,7 @@ describe("useInfiniteUsers", () => { }); it("should execute query for each admin role", async () => { - const adminRoles = [ - "Admin", - "Admin Viewer", - "proxy_admin", - "proxy_admin_viewer", - "org_admin", - ]; + const adminRoles = ["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer", "org_admin"]; for (const role of adminRoles) { vi.clearAllMocks(); @@ -328,12 +282,6 @@ describe("useInfiniteUsers", () => { expect(result.current.isSuccess).toBe(true); }); - expect(userListCall).toHaveBeenCalledWith( - "test-access-token", - null, - 1, - 50, - null, - ); + expect(userListCall).toHaveBeenCalledWith("test-access-token", null, 1, 50, null); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts index cb30299f46f..9031de3cb1a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts @@ -8,10 +8,7 @@ const infiniteUsersKeys = createQueryKeys("infiniteUsers"); const DEFAULT_PAGE_SIZE = 50; -export const useInfiniteUsers = ( - pageSize: number = DEFAULT_PAGE_SIZE, - searchEmail?: string, -) => { +export const useInfiniteUsers = (pageSize: number = DEFAULT_PAGE_SIZE, searchEmail?: string) => { const { accessToken, userRole } = useAuthorized(); return useInfiniteQuery({ queryKey: infiniteUsersKeys.list({ @@ -23,10 +20,10 @@ export const useInfiniteUsers = ( queryFn: async ({ pageParam }) => { return await userListCall( accessToken!, - null, // userIDs - pageParam as number, // page - pageSize, // page_size - searchEmail || null, // userEmail + null, // userIDs + pageParam as number, // page + pageSize, // page_size + searchEmail || null, // userEmail ); }, initialPageParam: 1, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 5bb55ee8d10..a611d619cc1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -69,17 +69,13 @@ function LayoutContent({ children }: { children: React.ReactNode }) { sidebarCollapsed={sidebarCollapsed} onToggleSidebar={toggleSidebar} proxySettings={undefined} - setProxySettings={() => { }} + setProxySettings={() => {}} accessToken={accessToken} />
- +
{children}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 944c56833e5..88f4382d7dd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -518,7 +518,11 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te ); } return ( - +
{visibleTabs.map((t) => t.tab)}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 32ab83ea754..045bf0a5f44 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -21,7 +21,7 @@ vi.mock("@/components/molecules/notifications_manager", () => ({ // Mock react-query const mockInvalidateQueries = vi.fn(); vi.mock("@tanstack/react-query", async (importOriginal) => { - const actual = await importOriginal() as any; + const actual = (await importOriginal()) as any; return { ...actual, useQueryClient: () => ({ @@ -178,24 +178,30 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-accessible", - model_info: { - id: "model-1", - access_via_team_ids: ["team-456"], - access_groups: [], + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-accessible", + model_info: { + id: "model-1", + access_via_team_ids: ["team-456"], + access_groups: [], + }, }, - }, - { - model_name: "gpt-3.5-turbo-blocked", - model_info: { - id: "model-2", - access_via_team_ids: ["team-789"], - access_groups: [], + { + model_name: "gpt-3.5-turbo-blocked", + model_info: { + id: "model-2", + access_via_team_ids: ["team-789"], + access_groups: [], + }, }, - }, - ], 2, 1, 1, 50); + ], + 2, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); @@ -239,24 +245,30 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-sales", - model_info: { - id: "model-sales-1", - access_via_team_ids: [], - access_groups: ["sales-model-group"], + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-sales", + model_info: { + id: "model-sales-1", + access_via_team_ids: [], + access_groups: ["sales-model-group"], + }, }, - }, - { - model_name: "gpt-4-engineering", - model_info: { - id: "model-eng-1", - access_via_team_ids: [], - access_groups: ["engineering-model-group"], + { + model_name: "gpt-4-engineering", + model_info: { + id: "model-eng-1", + access_via_team_ids: [], + access_groups: ["engineering-model-group"], + }, }, - }, - ], 2, 1, 1, 50); + ], + 2, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); @@ -284,26 +296,32 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-personal", - model_info: { - id: "model-personal-1", - direct_access: true, - access_via_team_ids: [], - access_groups: [], + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-personal", + model_info: { + id: "model-personal-1", + direct_access: true, + access_via_team_ids: [], + access_groups: [], + }, }, - }, - { - model_name: "gpt-4-team-only", - model_info: { - id: "model-team-1", - direct_access: false, - access_via_team_ids: ["team-123"], - access_groups: [], + { + model_name: "gpt-4-team-only", + model_info: { + id: "model-team-1", + direct_access: false, + access_via_team_ids: ["team-123"], + access_groups: [], + }, }, - }, - ], 2, 1, 1, 50); + ], + 2, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); @@ -330,38 +348,44 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-config", - litellm_model_name: "gpt-4-config", - provider: "openai", - model_info: { - id: "model-config-1", - db_model: false, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-config", + litellm_model_name: "gpt-4-config", + provider: "openai", + model_info: { + id: "model-config-1", + db_model: false, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, }, - }, - { - model_name: "gpt-4-db", - litellm_model_name: "gpt-4-db", - provider: "openai", - model_info: { - id: "model-db-1", - db_model: true, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", + { + model_name: "gpt-4-db", + litellm_model_name: "gpt-4-db", + provider: "openai", + model_info: { + id: "model-db-1", + db_model: true, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, }, - }, - ], 2, 1, 1, 50); + ], + 2, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); @@ -387,23 +411,29 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-config", - litellm_model_name: "gpt-4-config", - provider: "openai", - model_info: { - id: "model-config-1", - db_model: false, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-config", + litellm_model_name: "gpt-4-config", + provider: "openai", + model_info: { + id: "model-config-1", + db_model: false, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, }, - }, - ], 1, 1, 1, 50); + ], + 1, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); @@ -537,23 +567,29 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-delete-test", - litellm_model_name: "gpt-4-delete-test", - provider: "openai", - model_info: { - id: "model-to-delete", - db_model: true, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-delete-test", + litellm_model_name: "gpt-4-delete-test", + provider: "openai", + model_info: { + id: "model-to-delete", + db_model: true, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, }, - }, - ], 1, 1, 1, 50); + ], + 1, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); @@ -581,23 +617,29 @@ describe("AllModelsTab", () => { }), ); - const modelData = createPaginatedModelData([ - { - model_name: "gpt-4-clickable", - litellm_model_name: "gpt-4-clickable", - provider: "openai", - model_info: { - id: "clickable-model-id", - db_model: true, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", + const modelData = createPaginatedModelData( + [ + { + model_name: "gpt-4-clickable", + litellm_model_name: "gpt-4-clickable", + provider: "openai", + model_info: { + id: "clickable-model-id", + db_model: true, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, }, - }, - ], 1, 1, 1, 50); + ], + 1, + 1, + 1, + 50, + ); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 2626ace86d5..2aa1eb4c808 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -68,7 +68,7 @@ const AllModelsTab = ({ setCurrentPage(1); setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); }, 200), - [] + [], ); useEffect(() => { @@ -100,15 +100,11 @@ const AllModelsTab = ({ return sort.desc ? "desc" : "asc"; }, [sorting]); - const { data: rawModelData, isLoading: isLoadingModelsInfo, refetch: refetchModels } = useModelsInfo( - currentPage, - pageSize, - debouncedSearch || undefined, - undefined, - teamIdForQuery, - sortBy, - sortOrder - ); + const { + data: rawModelData, + isLoading: isLoadingModelsInfo, + refetch: refetchModels, + } = useModelsInfo(currentPage, pageSize, debouncedSearch || undefined, undefined, teamIdForQuery, sortBy, sortOrder); const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; const getProviderFromModel = (model: string) => { @@ -494,7 +490,7 @@ const AllModelsTab = ({ ) : ( {paginationMeta.total_count > 0 - ? `Showing ${((currentPage - 1) * pageSize) + 1} - ${Math.min(currentPage * pageSize, paginationMeta.total_count)} of ${paginationMeta.total_count} results` + ? `Showing ${(currentPage - 1) * pageSize + 1} - ${Math.min(currentPage * pageSize, paginationMeta.total_count)} of ${paginationMeta.total_count} results` : "Showing 0 results"} )} @@ -510,10 +506,9 @@ const AllModelsTab = ({ setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); }} disabled={currentPage === 1} - className={`px-3 py-1 text-sm border rounded-md ${currentPage === 1 - ? "bg-gray-100 text-gray-400 cursor-not-allowed" - : "hover:bg-gray-50" - }`} + className={`px-3 py-1 text-sm border rounded-md ${ + currentPage === 1 ? "bg-gray-100 text-gray-400 cursor-not-allowed" : "hover:bg-gray-50" + }`} > Previous @@ -529,10 +524,11 @@ const AllModelsTab = ({ setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); }} disabled={currentPage >= paginationMeta.total_pages} - className={`px-3 py-1 text-sm border rounded-md ${currentPage >= paginationMeta.total_pages - ? "bg-gray-100 text-gray-400 cursor-not-allowed" - : "hover:bg-gray-50" - }`} + className={`px-3 py-1 text-sm border rounded-md ${ + currentPage >= paginationMeta.total_pages + ? "bg-gray-100 text-gray-400 cursor-not-allowed" + : "hover:bg-gray-50" + }`} > Next @@ -550,8 +546,8 @@ const AllModelsTab = ({ setSelectedModelId, setSelectedTeamId, getDisplayModelName, - () => { }, - () => { }, + () => {}, + () => {}, expandedRows, setExpandedRows, setDeleteModalModelId, @@ -577,24 +573,28 @@ const AllModelsTab = ({ alertMessage="This action cannot be undone." message="Are you sure you want to delete this model?" resourceInformationTitle="Model Information" - resourceInformation={modelToDelete ? [ - { - label: "Model Name", - value: modelToDelete.model_name || "Not Set", - }, - { - label: "LiteLLM Model Name", - value: modelToDelete.litellm_model_name || "Not Set", - }, - { - label: "Provider", - value: modelToDelete.provider || "Not Set", - }, - { - label: "Created By", - value: modelToDelete.model_info?.created_by || "Not Set", - }, - ] : []} + resourceInformation={ + modelToDelete + ? [ + { + label: "Model Name", + value: modelToDelete.model_name || "Not Set", + }, + { + label: "LiteLLM Model Name", + value: modelToDelete.litellm_model_name || "Not Set", + }, + { + label: "Provider", + value: modelToDelete.provider || "Not Set", + }, + { + label: "Created By", + value: modelToDelete.model_info?.created_by || "Not Set", + }, + ] + : [] + } onCancel={() => setDeleteModalModelId(null)} onOk={handleDeleteModel} confirmLoading={deleteLoading} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx index 2b4b4ace491..abcbe80a382 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx @@ -36,43 +36,43 @@ export default function PlaygroundPage() { return (
- - - Chat - Compare - Compliance - Agent Builder (Experimental) - - - - - - - - - - - - - - - - + + + Chat + Compare + Compliance + Agent Builder (Experimental) + + + + + + + + + + + + + + + +
); } diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx index 25233725da6..6a61f5ed85f 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx @@ -304,8 +304,7 @@ describe("LoginPage", () => { }, writable: true, }); - document.cookie = - "token=; expires=Thu, 01 Jan 1970 00:00:00 GMT; path=/; SameSite=Lax"; + document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 GMT; path=/; SameSite=Lax"; }); afterEach(() => { diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 2db95947305..db3a069902a 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -50,13 +50,11 @@ function LoginPageContent() { // Validate the SSO code is a plausible OAuth authorization code (alphanumeric // plus common URL-safe chars) so that arbitrary user input cannot trigger the // exchange endpoint. - const ssoCode = - rawSsoCode && /^[a-zA-Z0-9._~+/=-]+$/.test(rawSsoCode) ? rawSsoCode : null; + const ssoCode = rawSsoCode && /^[a-zA-Z0-9._~+/=-]+$/.test(rawSsoCode) ? rawSsoCode : null; if (ssoCode) { const rawWorkerUrl = localStorage.getItem("litellm_worker_url"); // Validate the stored worker URL: only allow http(s) URLs. - const workerUrl = - rawWorkerUrl && /^https?:\/\/.+/.test(rawWorkerUrl) ? rawWorkerUrl : null; + const workerUrl = rawWorkerUrl && /^https?:\/\/.+/.test(rawWorkerUrl) ? rawWorkerUrl : null; exchangeLoginCode(ssoCode, workerUrl).then(() => { params.delete("code"); const cleanSearch = params.toString(); @@ -277,10 +275,7 @@ function LoginPageContent() { {!uiConfig?.sso_configured ? ( - + @@ -315,7 +310,13 @@ function LoginPageContent() { type="info" showIcon closable - message={Single Sign-On (SSO) is enabled. LiteLLM no longer automatically redirects to the SSO login flow upon loading this page. To re-enable auto-redirect-to-SSO, set AUTO_REDIRECT_UI_LOGIN_TO_SSO=true in your environment configuration.} + message={ + + Single Sign-On (SSO) is enabled. LiteLLM no longer automatically redirects to the SSO login flow upon + loading this page. To re-enable auto-redirect-to-SSO, set{" "} + AUTO_REDIRECT_UI_LOGIN_TO_SSO=true in your environment configuration. + + } /> )} diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index 3b3729c1ac9..d292925f810 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -80,12 +80,12 @@ const McpOAuthCallbackContent = () => {

LiteLLM MCP OAuth

-

- Authorization complete. You may close this window and return to the LiteLLM dashboard. -

-

- If the window does not close automatically, everything is still saved—you can close it manually. -

+

+ Authorization complete. You may close this window and return to the LiteLLM dashboard. +

+

+ If the window does not close automatically, everything is still saved—you can close it manually. +

); diff --git a/ui/litellm-dashboard/src/app/model_hub_table/page.tsx b/ui/litellm-dashboard/src/app/model_hub_table/page.tsx index f35a6943a63..472cea4c27e 100644 --- a/ui/litellm-dashboard/src/app/model_hub_table/page.tsx +++ b/ui/litellm-dashboard/src/app/model_hub_table/page.tsx @@ -16,9 +16,7 @@ function PublicModelHubTableContent() { setAccessToken(key); }, [key]); - return ( - - ); + return ; } export default function PublicModelHubTable() { diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx index d7a7ffb1b15..59071f17bf6 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx @@ -11,9 +11,7 @@ describe("OnboardingErrorView", () => { it("should show the expiry description", () => { render(); - expect( - screen.getByText("The invitation link may be invalid or expired.") - ).toBeInTheDocument(); + expect(screen.getByText("The invitation link may be invalid or expired.")).toBeInTheDocument(); }); it("should render a Back to Login link pointing to /ui/login", () => { diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx index add58102c57..23a8bc6725a 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx @@ -26,9 +26,7 @@ export function OnboardingForm({ variant }: OnboardingFormProps) { const { mutate: claimToken, isPending } = useClaimOnboardingToken(); - const decoded = credentialsData?.token - ? (jwtDecode(credentialsData.token) as { [key: string]: any }) - : null; + const decoded = credentialsData?.token ? (jwtDecode(credentialsData.token) as { [key: string]: any }) : null; const userEmail: string = decoded?.user_email ?? ""; const userId: string | null = decoded?.user_id ?? null; const accessToken: string | null = decoded?.key ?? null; @@ -53,14 +51,12 @@ export function OnboardingForm({ variant }: OnboardingFormProps) { clearTokenCookies(); storeLoginToken(data.token); const proxyBaseUrl = getProxyBaseUrl(); - window.location.href = proxyBaseUrl - ? `${proxyBaseUrl}/ui/?login=success` - : "/ui/?login=success"; + window.location.href = proxyBaseUrl ? `${proxyBaseUrl}/ui/?login=success` : "/ui/?login=success"; }, onError: (error: Error) => { setClaimError(error.message || "Failed to submit. Please try again."); }, - } + }, ); }; diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.test.tsx index f742176d1ba..f3286984706 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.test.tsx @@ -74,16 +74,12 @@ describe("OnboardingFormBody", () => { await user.click(screen.getByRole("button", { name: /sign up/i })); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith( - expect.objectContaining({ password: "mypassword" }) - ); + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ password: "mypassword" })); }); }); it("should show 'Reset Password' on the submit button for reset_password variant", () => { render(); - expect( - screen.getByRole("button", { name: /reset password/i }) - ).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /reset password/i })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx index c57c7328b61..4aa5e2e6138 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx @@ -9,13 +9,7 @@ type OnboardingFormBodyProps = { onSubmit: (values: { password: string }) => void; }; -export function OnboardingFormBody({ - variant, - userEmail, - isPending, - claimError, - onSubmit, -}: OnboardingFormBodyProps) { +export function OnboardingFormBody({ variant, userEmail, isPending, claimError, onSubmit }: OnboardingFormBodyProps) { const [form] = Form.useForm(); React.useEffect(() => { @@ -28,9 +22,7 @@ export function OnboardingFormBody({ 🚅 LiteLLM - - {variant === "reset_password" ? "Reset Password" : "Sign Up"} - + {variant === "reset_password" ? "Reset Password" : "Sign Up"} {variant === "reset_password" ? "Reset your password to access Admin UI." @@ -45,12 +37,7 @@ export function OnboardingFormBody({ description={
SSO is under the Enterprise Tier. -
@@ -59,7 +46,12 @@ export function OnboardingFormBody({ /> )} -
onSubmit({ password: values.password })}> + onSubmit({ password: values.password })} + > @@ -68,18 +60,12 @@ export function OnboardingFormBody({ label="Password" name="password" rules={[{ required: true, message: "password required to sign up" }]} - help={ - variant === "reset_password" - ? "Enter your new password" - : "Create a password for your account" - } + help={variant === "reset_password" ? "Enter your new password" : "Create a password for your account"} > - {claimError && ( - - )} + {claimError && }
- } - > + Loading...}> ); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index ce12967c911..da6a0d5a76f 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -46,7 +46,13 @@ import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; import { useAuth } from "@/contexts/AuthContext"; -import { buildLoginUrlWithReturn, consumeReturnUrl, isValidReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; +import { + buildLoginUrlWithReturn, + consumeReturnUrl, + isValidReturnUrl, + normalizeUrlForCompare, + storeReturnUrl, +} from "@/utils/returnUrlUtils"; import { isAdminRole } from "@/utils/roles"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { useRouter, useSearchParams } from "next/navigation"; @@ -67,17 +73,8 @@ interface ProxySettings { const LEGACY_REDIRECTS: Record = {}; function CreateKeyPageContent() { - const { - authLoading, - token, - userID, - userRole, - userEmail, - accessToken, - premiumUser, - setUserRole, - setUserEmail, - } = useAuth(); + const { authLoading, token, userID, userRole, userEmail, accessToken, premiumUser, setUserRole, setUserEmail } = + useAuth(); const [teams, setTeams] = useState(null); const [keys, setKeys] = useState([]); @@ -129,15 +126,13 @@ function CreateKeyPageContent() { // Validate owned_by against allowed values const validOwnedByValues = ["you", "service_account", "another_user"]; - const validatedOwnedBy = ownedBy && validOwnedByValues.includes(ownedBy) - ? (ownedBy as CreateKeyPrefillData["owned_by"]) - : undefined; + const validatedOwnedBy = + ownedBy && validOwnedByValues.includes(ownedBy) ? (ownedBy as CreateKeyPrefillData["owned_by"]) : undefined; // Validate key_type against allowed values const validKeyTypes = ["default", "llm_api", "management"]; - const validatedKeyType = keyType && validKeyTypes.includes(keyType) - ? (keyType as CreateKeyPrefillData["key_type"]) - : undefined; + const validatedKeyType = + keyType && validKeyTypes.includes(keyType) ? (keyType as CreateKeyPrefillData["key_type"]) : undefined; // Sanitize key_alias (limit length, trim whitespace) const sanitizedKeyAlias = keyAlias @@ -149,8 +144,8 @@ function CreateKeyPageContent() { ? modelsParam .split(",") .slice(0, 100) // Limit number of models to prevent DoS - .map(m => m.trim().slice(0, 256)) // Limit individual model name length - .filter(m => m.length > 0) // Remove empty strings + .map((m) => m.trim().slice(0, 256)) // Limit individual model name length + .filter((m) => m.length > 0) // Remove empty strings : undefined; return { @@ -259,7 +254,9 @@ function CreateKeyPageContent() { if (accessToken && userID && userRole) { v2TeamListCall(accessToken, 1, 100, { userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null, - }).then((response) => setTeams(response.teams ?? [])).catch(console.error); + }) + .then((response) => setTeams(response.teams ?? [])) + .catch(console.error); } if (accessToken) { fetchOrganizations(accessToken, setOrganizations); @@ -353,235 +350,231 @@ function CreateKeyPageContent() { return ( }> - - - {invitation_id ? ( - + + {invitation_id ? ( + + ) : ( +
+ - ) : ( -
- -
-
+
+
- {page == "api-keys" ? ( - + ) : page == "models" ? ( + + ) : page == "llm-playground" ? ( + + ) : page == "users" ? ( + + ) : page == "teams" ? ( + + ) : page == "organizations" ? ( + + ) : page == "admin-panel" ? ( + + ) : page == "api_ref" || page == "api-reference" ? ( + + ) : page == "logging-and-alerts" ? ( + + ) : page == "budgets" ? ( + + ) : page == "guardrails" ? ( + + ) : page == "policies" ? ( + + ) : page == "agents" ? ( + + ) : page == "prompts" ? ( + + ) : page == "transform-request" ? ( + + ) : page == "router-settings" ? ( + + ) : page == "ui-theme" ? ( + + ) : page == "cost-tracking" ? ( + + ) : page == "model-hub-table" ? ( + isAdminRole(userRole) ? ( + - ) : page == "models" ? ( - - ) : page == "llm-playground" ? ( - - ) : page == "users" ? ( - - ) : page == "teams" ? ( - - ) : page == "organizations" ? ( - - ) : page == "admin-panel" ? ( - - ) : page == "api_ref" || page == "api-reference" ? ( - - ) : page == "logging-and-alerts" ? ( - - ) : page == "budgets" ? ( - - ) : page == "guardrails" ? ( - - ) : page == "policies" ? ( - - ) : page == "agents" ? ( - - ) : page == "prompts" ? ( - - ) : page == "transform-request" ? ( - - ) : page == "router-settings" ? ( - - ) : page == "ui-theme" ? ( - - ) : page == "cost-tracking" ? ( - - ) : page == "model-hub-table" ? ( - isAdminRole(userRole) ? ( - - ) : ( - - ) - ) : page == "caching" ? ( - - ) : page == "pass-through-settings" ? ( - - ) : page == "logs" ? ( - - ) : page == "mcp-servers" ? ( - - ) : page == "search-tools" ? ( - - ) : page == "tag-management" ? ( - - ) : page == "skills" || page == "claude-code-plugins" ? ( - - ) : page == "access-groups" ? ( - - ) : page == "projects" ? ( - - ) : page == "vector-stores" ? ( - - ) : page == "tool-policies" ? ( - - ) : page == "workflows" ? ( - - ) : page == "memory" ? ( - - ) : page == "guardrails-monitor" ? ( - - ) : page == "new_usage" ? ( - ) : ( - - )} -
- - {/* Survey Components */} - - - - {/* Claude Code Components */} - - + + ) + ) : page == "caching" ? ( + + ) : page == "pass-through-settings" ? ( + + ) : page == "logs" ? ( + + ) : page == "mcp-servers" ? ( + + ) : page == "search-tools" ? ( + + ) : page == "tag-management" ? ( + + ) : page == "skills" || page == "claude-code-plugins" ? ( + + ) : page == "access-groups" ? ( + + ) : page == "projects" ? ( + + ) : page == "vector-stores" ? ( + + ) : page == "tool-policies" ? ( + + ) : page == "workflows" ? ( + + ) : page == "memory" ? ( + + ) : page == "guardrails-monitor" ? ( + + ) : page == "new_usage" ? ( + + ) : ( + + )}
- )} - - + + {/* Survey Components */} + + + + {/* Claude Code Components */} + + +
+ )} + + ); } diff --git a/ui/litellm-dashboard/src/components/AIHub/AgentHubTableColumns.test.tsx b/ui/litellm-dashboard/src/components/AIHub/AgentHubTableColumns.test.tsx index 083e67c297a..f980aee3c2c 100644 --- a/ui/litellm-dashboard/src/components/AIHub/AgentHubTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/AgentHubTableColumns.test.tsx @@ -104,12 +104,10 @@ describe("AgentHubTableColumns", () => { render(); // "In:" and "Out:" are in children; getByText with exact:false // matches against the element's full textContent across child nodes - expect(screen.getByText((_, el) => - el?.tagName === "P" && el.textContent === "In: text" - )).toBeInTheDocument(); - expect(screen.getByText((_, el) => - el?.tagName === "P" && el.textContent === "Out: text, image" - )).toBeInTheDocument(); + expect(screen.getByText((_, el) => el?.tagName === "P" && el.textContent === "In: text")).toBeInTheDocument(); + expect( + screen.getByText((_, el) => el?.tagName === "P" && el.textContent === "Out: text, image"), + ).toBeInTheDocument(); }); it("should display 'Yes' badge for public agents", () => { diff --git a/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx b/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx index 043077c0210..762b0836921 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx @@ -2,14 +2,8 @@ import { SearchOutlined } from "@ant-design/icons"; import { Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; import { Input } from "antd"; import React, { useEffect, useMemo, useState } from "react"; -import { - extractCategories, - filterPluginsByCategory, - filterPluginsBySearch, -} from "../claude_code_plugins/helpers"; -import { - MarketplaceResponse -} from "../claude_code_plugins/types"; +import { extractCategories, filterPluginsByCategory, filterPluginsBySearch } from "../claude_code_plugins/helpers"; +import { MarketplaceResponse } from "../claude_code_plugins/types"; import { ModelDataTable } from "../model_dashboard/table"; import NotificationsManager from "../molecules/notifications_manager"; import { getClaudeCodeMarketplace } from "../networking"; @@ -19,11 +13,8 @@ interface ClaudeCodeMarketplaceTabProps { publicPage?: boolean; } -const ClaudeCodeMarketplaceTab: React.FC = ({ - publicPage = false, -}) => { - const [marketplaceData, setMarketplaceData] = - useState(null); +const ClaudeCodeMarketplaceTab: React.FC = ({ publicPage = false }) => { + const [marketplaceData, setMarketplaceData] = useState(null); const [isLoading, setIsLoading] = useState(true); const [searchTerm, setSearchTerm] = useState(""); const [selectedCategoryIndex, setSelectedCategoryIndex] = useState(0); @@ -74,18 +65,13 @@ const ClaudeCodeMarketplaceTab: React.FC = ({ return plugins; }, [marketplaceData, selectedCategory, searchTerm]); - const columns = useMemo( - () => getMarketplaceTableColumns(copyToClipboard, publicPage), - [publicPage] - ); + const columns = useMemo(() => getMarketplaceTableColumns(copyToClipboard, publicPage), [publicPage]); if (!marketplaceData && !isLoading) { return (
- - Failed to load marketplace. Please try again later. - + Failed to load marketplace. Please try again later.
); @@ -110,14 +96,8 @@ const ClaudeCodeMarketplaceTab: React.FC = ({ {categories.map((category) => { // Count plugins in this category - const categoryPlugins = filterPluginsByCategory( - marketplaceData?.plugins || [], - category - ); - const count = filterPluginsBySearch( - categoryPlugins, - searchTerm - ).length; + const categoryPlugins = filterPluginsByCategory(marketplaceData?.plugins || [], category); + const count = filterPluginsBySearch(categoryPlugins, searchTerm).length; return ( @@ -143,8 +123,7 @@ const ClaudeCodeMarketplaceTab: React.FC = ({ {/* Footer Info */}
- Showing {filteredPlugins.length} of{" "} - {marketplaceData?.plugins.length || 0} plugin + Showing {filteredPlugins.length} of {marketplaceData?.plugins.length || 0} plugin {marketplaceData?.plugins.length !== 1 ? "s" : ""} {searchTerm && ` matching "${searchTerm}"`} {selectedCategory !== "All" && ` in ${selectedCategory}`} diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index ee59ac84ece..3a22a55298e 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -48,11 +48,7 @@ describe("ModelHubTable", () => { }); // Reusable helper function to setup mocks for auth redirect tests - const setupAuthRedirectTest = ( - requireAuth: boolean, - tokenValue: string | null, - isTokenValid: boolean - ) => { + const setupAuthRedirectTest = (requireAuth: boolean, tokenValue: string | null, isTokenValid: boolean) => { mockUseUISettings.mockReturnValue({ data: { values: { @@ -87,14 +83,12 @@ describe("ModelHubTable", () => { tokenValue: string | null, isTokenValid: boolean, shouldRedirect: boolean, - description: string + description: string, ) => { it(description, async () => { setupAuthRedirectTest(requireAuth, tokenValue, isTokenValid); - renderWithProviders( - - ); + renderWithProviders(); await waitFor(() => { if (shouldRedirect) { @@ -125,7 +119,9 @@ describe("ModelHubTable", () => { isLoading: false, }); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByText("AI Hub")).toBeInTheDocument(); @@ -172,7 +168,7 @@ describe("ModelHubTable", () => { null, false, true, - "should redirect to login when requireAuth is true and there is no token" + "should redirect to login when requireAuth is true and there is no token", ); testAuthRedirect( @@ -180,7 +176,7 @@ describe("ModelHubTable", () => { "expired-token", false, true, - "should redirect to login when requireAuth is true and token is expired" + "should redirect to login when requireAuth is true and token is expired", ); testAuthRedirect( @@ -188,24 +184,18 @@ describe("ModelHubTable", () => { "malformed-token", false, true, - "should redirect to login when requireAuth is true and token is malformed" + "should redirect to login when requireAuth is true and token is malformed", ); // Test cases where requireAuth is false - should NOT redirect regardless of token state - testAuthRedirect( - false, - null, - false, - false, - "should not redirect when requireAuth is false and there is no token" - ); + testAuthRedirect(false, null, false, false, "should not redirect when requireAuth is false and there is no token"); testAuthRedirect( false, "expired-token", false, false, - "should not redirect when requireAuth is false and token is expired" + "should not redirect when requireAuth is false and token is expired", ); testAuthRedirect( @@ -213,7 +203,7 @@ describe("ModelHubTable", () => { "malformed-token", false, false, - "should not redirect when requireAuth is false and token is malformed" + "should not redirect when requireAuth is false and token is malformed", ); }); }); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 75058157a65..5d171139ab5 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -526,9 +526,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {publicPage == false && canModify && (
- +
)} = ({ s.description?.toLowerCase().includes(q) || s.domain?.toLowerCase().includes(q) || s.namespace?.toLowerCase().includes(q) || - s.keywords?.some((k) => k.toLowerCase().includes(q)) + s.keywords?.some((k) => k.toLowerCase().includes(q)), ); } return result; @@ -94,9 +94,7 @@ const SkillHubDashboard: React.FC = ({ {/* Search + filters + table */}
-

- All {publicPage ? "Public " : ""}Skills -

+

All {publicPage ? "Public " : ""}Skills

+ - -