From ac5c5ee302729a9833277e99b74afbef7948b51e Mon Sep 17 00:00:00 2001 From: Quentin Machu Date: Mon, 26 Jan 2026 17:42:56 -0500 Subject: [PATCH 01/13] fix: Search tools not found when using per-request routers **Root Cause:** When API keys or teams have router_settings configured, the proxy creates per-request Router instances from user_config. These new routers were missing search_tools from the main router, causing "search tool not found" errors despite search_tools being configured. **The Fix:** 1. **common_request_processing.py (lines 554-556):** Pass search_tools from main router to user_config so per-request routers inherit them 2. **proxy_server.py (line 3171):** Remove `if len(_model_list) > 0` check to allow router creation with empty model list (needed for search-tools-only use case) 3. **proxy_server.py (line 746):** Remove redundant search_tools loading code (already handled by _init_search_tools_in_db() called during startup) --- litellm/proxy/common_request_processing.py | 7 ++++++- litellm/proxy/proxy_server.py | 23 +++++++++++----------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0d3e61b75c7..4f0c59474f1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -547,9 +547,14 @@ class ProxyBaseLLMRequestProcessing: # Get model_list from current router model_list = llm_router.get_model_list() if model_list is not None: - # Create user_config with model_list and router_settings + # Create user_config with model_list, search_tools, and router_settings # This creates a per-request router with the hierarchical settings user_config = {"model_list": model_list, **router_settings} + + # Include search_tools from main router so per-request router has them + if hasattr(llm_router, "search_tools") and llm_router.search_tools: + user_config["search_tools"] = llm_router.search_tools + self.data["user_config"] = user_config if "messages" in self.data and self.data["messages"]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 16cdd9da64b..056367243f9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3173,17 +3173,18 @@ class ProxyConfig: _model_list: list = self.decrypt_model_list_from_db( new_models=models_list ) - if len(_model_list) > 0: - verbose_proxy_logger.debug(f"_model_list: {_model_list}") - llm_router = litellm.Router( - model_list=_model_list, - router_general_settings=RouterGeneralSettings( - async_only_mode=True # only init async clients - ), - search_tools=search_tools, - ignore_invalid_deployments=True, - ) - verbose_proxy_logger.debug(f"updated llm_router: {llm_router}") + # Create router even with empty model list to support search_tools + # Router can function with model_list=[] and only search_tools + verbose_proxy_logger.debug(f"_model_list: {_model_list}") + llm_router = litellm.Router( + model_list=_model_list, + router_general_settings=RouterGeneralSettings( + async_only_mode=True # only init async clients + ), + search_tools=search_tools, + ignore_invalid_deployments=True, + ) + verbose_proxy_logger.debug(f"updated llm_router: {llm_router}") else: verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") ## DELETE MODEL LOGIC From a9eae5937fb58a53047e581d4950b10b956286ce Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 31 Jan 2026 16:04:52 -0800 Subject: [PATCH 02/13] Override router settings --- litellm/proxy/common_request_processing.py | 16 +-- litellm/proxy/route_llm_request.py | 35 +++-- litellm/router.py | 11 +- .../proxy/test_common_request_processing.py | 34 +++-- .../proxy/test_route_llm_request.py | 131 ++++++++++++------ 5 files changed, 144 insertions(+), 83 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b69f175e5d2..ee38d56c855 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -626,20 +626,10 @@ class ProxyBaseLLMRequestProcessing: ) # If router_settings found (from key, team, or global), apply them - # This ensures key/team settings override global settings + # Pass settings as per-request overrides instead of creating a new Router + # This avoids expensive Router instantiation on each request if router_settings is not None and router_settings: - # Get model_list from current router - model_list = llm_router.get_model_list() - if model_list is not None: - # Create user_config with model_list, search_tools, and router_settings - # This creates a per-request router with the hierarchical settings - user_config = {"model_list": model_list, **router_settings} - - # Include search_tools from main router so per-request router has them - if hasattr(llm_router, "search_tools") and llm_router.search_tools: - user_config["search_tools"] = llm_router.search_tools - - self.data["user_config"] = user_config + self.data["router_settings_override"] = router_settings if "messages" in self.data and self.data["messages"]: logging_obj.update_messages(self.data["messages"]) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index c6a93164d49..e2749eb8187 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -197,18 +197,33 @@ async def route_request( models = [model.strip() for model in data.pop("model").split(",")] return llm_router.abatch_completion(models=models, **data) - elif "user_config" in data: - router_config = data.pop("user_config") + elif "router_settings_override" in data: + # Apply per-request router settings overrides from key/team config + # Instead of creating a new Router (expensive), merge settings into kwargs + # The Router already supports per-request overrides for these settings + override_settings = data.pop("router_settings_override") - # Filter router_config to only include valid Router.__init__ arguments - # This prevents TypeError when invalid parameters are stored in the database - valid_args = litellm.Router.get_valid_args() - filtered_config = {k: v for k, v in router_config.items() if k in valid_args} + # Settings that the Router accepts as per-request kwargs + # These override the global router settings for this specific request + per_request_settings = [ + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "timeout", + "model_group_retry_policy", + ] - user_router = litellm.Router(**filtered_config) - ret_val = getattr(user_router, f"{route_type}")(**data) - user_router.discard() - return ret_val + # Merge override settings into data (only if not already set in request) + for key in per_request_settings: + if key in override_settings and key not in data: + data[key] = override_settings[key] + + # Use main router with overridden kwargs + if llm_router is not None: + return getattr(llm_router, f"{route_type}")(**data) + else: + return getattr(litellm, f"{route_type}")(**data) elif llm_router is not None: # Skip model-based routing for container operations if route_type in [ diff --git a/litellm/router.py b/litellm/router.py index 6c191c8ab03..ed480d6468a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4880,6 +4880,10 @@ class Router: content_policy_fallbacks = kwargs.pop( "content_policy_fallbacks", self.content_policy_fallbacks ) + # Support per-request model_group_retry_policy override (from key/team settings) + model_group_retry_policy = kwargs.pop( + "model_group_retry_policy", self.model_group_retry_policy + ) model_group: Optional[str] = kwargs.get("model") num_retries = kwargs.pop("num_retries") @@ -4928,7 +4932,7 @@ class Router: _retry_policy_applies = False if ( self.retry_policy is not None - or self.model_group_retry_policy is not None + or model_group_retry_policy is not None ): # get num_retries from retry policy # Use the model_group captured at the start of the function, or get it from metadata @@ -4936,9 +4940,12 @@ class Router: _model_group_for_retry_policy = ( model_group or _metadata.get("model_group") or kwargs.get("model") ) - _retry_policy_retries = self.get_num_retries_from_retry_policy( + # Use per-request model_group_retry_policy if provided, otherwise use self + _retry_policy_retries = _get_num_retries_from_retry_policy( exception=original_exception, model_group=_model_group_for_retry_policy, + model_group_retry_policy=model_group_retry_policy, + retry_policy=self.retry_policy, ) if _retry_policy_retries is not None: num_retries = _retry_policy_retries diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 6edcdab15c0..d2abe80977a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -78,9 +78,16 @@ class TestProxyBaseLLMRequestProcessing: assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] @pytest.mark.asyncio - async def test_should_apply_hierarchical_router_settings_to_user_config( + async def test_should_apply_hierarchical_router_settings_as_override( self, monkeypatch ): + """ + Test that hierarchical router settings are stored as router_settings_override + instead of creating a full user_config with model_list. + + This approach avoids expensive per-request Router instantiation by passing + settings as kwargs overrides to the main router. + """ processing_obj = ProxyBaseLLMRequestProcessing(data={}) mock_request = MagicMock(spec=Request) mock_request.headers = {} @@ -117,12 +124,7 @@ class TestProxyBaseLLMRequestProcessing: return_value=mock_router_settings ) - mock_model_list = [ - {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo"}}, - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}, - ] mock_llm_router = MagicMock() - mock_llm_router.get_model_list = MagicMock(return_value=mock_model_list) mock_prisma_client = MagicMock() monkeypatch.setattr( @@ -146,14 +148,20 @@ class TestProxyBaseLLMRequestProcessing: user_api_key_dict=mock_user_api_key_dict, prisma_client=mock_prisma_client, ) - mock_llm_router.get_model_list.assert_called_once() + # get_model_list should NOT be called - we no longer copy model list for per-request routers + mock_llm_router.get_model_list.assert_not_called() - assert "user_config" in returned_data - user_config = returned_data["user_config"] - assert user_config["model_list"] == mock_model_list - assert user_config["routing_strategy"] == "least-busy" - assert user_config["timeout"] == 30.0 - assert user_config["num_retries"] == 3 + # Settings should be stored as router_settings_override (not user_config) + # This allows passing them as kwargs to the main router instead of creating a new one + assert "router_settings_override" in returned_data + assert "user_config" not in returned_data + + router_settings_override = returned_data["router_settings_override"] + assert router_settings_override["routing_strategy"] == "least-busy" + assert router_settings_override["timeout"] == 30.0 + assert router_settings_override["num_retries"] == 3 + # model_list should NOT be in the override settings + assert "model_list" not in router_settings_override @pytest.mark.asyncio async def test_stream_timeout_header_processing(self): diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 90eace63714..1283d2ccbe7 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -137,62 +137,103 @@ async def test_route_request_no_model_required_with_router_settings_and_no_route @pytest.mark.asyncio -async def test_route_request_with_invalid_router_params(): +async def test_route_request_with_router_settings_override(): """ - Test that route_request filters out invalid Router init params from 'user_config'. - This covers the fix for https://github.com/BerriAI/litellm/issues/19693 + Test that route_request handles router_settings_override by merging settings into kwargs + instead of creating a new Router (which is expensive and was the old behavior). """ - import litellm - from litellm.router import Router - from unittest.mock import AsyncMock - - # Mock data with user_config containing invalid keys (simulating DB entry) + # Mock data with router_settings_override containing per-request settings data = { "model": "gpt-3.5-turbo", - "user_config": { - "model_list": [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "test"}, - } - ], - "model_alias_map": {"alias": "real_model"}, # INVALID PARAM - "invalid_garbage_key": "crash_me", # INVALID PARAM + "messages": [{"role": "user", "content": "Hello"}], + "router_settings_override": { + "fallbacks": [{"gpt-3.5-turbo": ["gpt-4"]}], + "num_retries": 5, + "timeout": 30, + "model_group_retry_policy": {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}}, + # These settings should be ignored (not in per_request_settings list) + "routing_strategy": "least-busy", + "model_group_alias": {"alias": "real_model"}, }, } - # We expect Router(**config) to succeed because of the filtering. - # If filtering fails, this will raise TypeError and fail the test. + llm_router = MagicMock() + llm_router.acompletion.return_value = "success" + + response = await route_request(data, llm_router, None, "acompletion") + + assert response == "success" + # Verify the router method was called with merged settings + call_kwargs = llm_router.acompletion.call_args[1] + assert call_kwargs["fallbacks"] == [{"gpt-3.5-turbo": ["gpt-4"]}] + assert call_kwargs["num_retries"] == 5 + assert call_kwargs["timeout"] == 30 + assert call_kwargs["model_group_retry_policy"] == {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}} + # Verify unsupported settings were NOT merged + assert "routing_strategy" not in call_kwargs + assert "model_group_alias" not in call_kwargs + # Verify router_settings_override was removed from data + assert "router_settings_override" not in call_kwargs + + +@pytest.mark.asyncio +async def test_route_request_with_router_settings_override_no_router(): + """ + Test that router_settings_override works when no router is provided, + falling back to litellm module directly. + """ + import litellm + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "router_settings_override": { + "fallbacks": [{"gpt-3.5-turbo": ["gpt-4"]}], + "num_retries": 3, + }, + } + + # Use MagicMock explicitly to avoid auto-AsyncMock behavior in Python 3.12+ + mock_completion = MagicMock(return_value="success") + original_acompletion = litellm.acompletion + litellm.acompletion = mock_completion + try: - # route_request calls getattr(user_router, route_type)(**data) - # We'll mock the internal call to avoid making real network requests - with pytest.MonkeyPatch.context() as m: - # Mock the method that gets called on the router instance - # We don't easily have access to the instance created INSIDE existing route_request - # So we will wrap litellm.Router to spy on it or verify it doesn't crash + response = await route_request(data, None, None, "acompletion") - original_router_init = litellm.Router.__init__ + assert response == "success" + # Verify litellm.acompletion was called with merged settings + call_kwargs = mock_completion.call_args[1] + assert call_kwargs["fallbacks"] == [{"gpt-3.5-turbo": ["gpt-4"]}] + assert call_kwargs["num_retries"] == 3 + finally: + litellm.acompletion = original_acompletion - def safe_router_init(self, **kwargs): - # Verify that invalid keys are NOT present in kwargs - assert "model_alias_map" not in kwargs - assert "invalid_garbage_key" not in kwargs - # Call original init (which would raise TypeError if invalid keys were present) - original_router_init(self, **kwargs) - m.setattr(litellm.Router, "__init__", safe_router_init) +@pytest.mark.asyncio +async def test_route_request_with_router_settings_override_preserves_existing(): + """ + Test that router_settings_override does not override settings already in the request. + Request-level settings take precedence over key/team settings. + """ + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "num_retries": 10, # Request-level setting + "router_settings_override": { + "num_retries": 3, # Key/team setting - should NOT override + "timeout": 30, # Key/team setting - should be applied + }, + } - # Use 'acompletion' as the route_type - # We also need to mock the completion method to avoid real calls - m.setattr(Router, "acompletion", AsyncMock(return_value="success")) + llm_router = MagicMock() + llm_router.acompletion.return_value = "success" - response = await route_request(data, None, None, "acompletion") - assert response == "success" + response = await route_request(data, llm_router, None, "acompletion") - except TypeError as e: - pytest.fail( - f"route_request raised TypeError, implying invalid params were passed to Router: {e}" - ) - except Exception: - # Other exceptions might happen (e.g. valid config issues) but we care about TypeError here - pass + assert response == "success" + call_kwargs = llm_router.acompletion.call_args[1] + # Request-level num_retries should take precedence + assert call_kwargs["num_retries"] == 10 + # Key/team timeout should be applied since not in request + assert call_kwargs["timeout"] == 30 From 08a6fe2bfa2ad937bc2e340ab47b35399857fed2 Mon Sep 17 00:00:00 2001 From: Swayambhu Date: Fri, 6 Feb 2026 09:44:02 +0530 Subject: [PATCH 03/13] refactor: migrate Ant Design notifications to use `App.useApp()` context via a new global provider. 1. notifications_manager.tsx - Hybrid notification approach: Added notificationInstance variable to store the context-based instance Added setNotificationInstance() function to inject the instance from context Created getNotification() helper that prefers context instance, falls back to static Added COMMON_NOTIFICATION_PROPS (exported) with showProgress: true and pauseOnHover: true All notification methods ( error , warning , info , success , fromBackend ) now spread COMMON_NOTIFICATION_PROPS 2. AntdGlobalProvider.tsx - New context provider: Wraps app with Antd's component Uses App.useApp() hook to get the context-based notification instance Injects it into NotificationManager via setNotificationInstance() --- ui/litellm-dashboard/src/app/layout.tsx | 6 ++- .../molecules/notifications_manager.tsx | 38 ++++++++++++------- .../src/contexts/AntdGlobalProvider.tsx | 25 ++++++++++++ 3 files changed, 54 insertions(+), 15 deletions(-) create mode 100644 ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx diff --git a/ui/litellm-dashboard/src/app/layout.tsx b/ui/litellm-dashboard/src/app/layout.tsx index 95c485fe2f0..1233da9046f 100644 --- a/ui/litellm-dashboard/src/app/layout.tsx +++ b/ui/litellm-dashboard/src/app/layout.tsx @@ -2,6 +2,8 @@ import type { Metadata } from "next"; import { Inter } from "next/font/google"; import "./globals.css"; +import AntdGlobalProvider from "@/contexts/AntdGlobalProvider"; + const inter = Inter({ subsets: ["latin"] }); export const metadata: Metadata = { @@ -17,7 +19,9 @@ export default function RootLayout({ }>) { return ( - {children} + + {children} + ); } diff --git a/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx b/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx index 71b14bee885..59b048b412c 100644 --- a/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx +++ b/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx @@ -1,8 +1,18 @@ import React from "react"; -import { notification } from "antd"; +import { notification as staticNotification } from "antd"; +import type { NotificationInstance } from "antd/es/notification/interface"; import { parseErrorMessage } from "../shared/errorUtils"; import { ArgsProps } from "antd/es/notification"; +let notificationInstance: NotificationInstance | null = null; + +export const setNotificationInstance = (instance: NotificationInstance) => { + notificationInstance = instance; +}; + +// Helper to get the best available notification instance +const getNotification = () => notificationInstance || staticNotification; + type Placement = "top" | "topLeft" | "topRight" | "bottom" | "bottomLeft" | "bottomRight"; type NotificationConfig = { @@ -251,7 +261,7 @@ function looksErrorPayload(input: any, status?: number): boolean { const NotificationManager = { error(input: string | NotificationConfig) { const cfg = normalize(input, "Error"); - notification.error({ + getNotification().error({ ...COMMON_NOTIFICATION_PROPS, ...cfg, placement: cfg.placement ?? defaultPlacement(), @@ -261,7 +271,7 @@ const NotificationManager = { warning(input: string | NotificationConfig) { const cfg = normalize(input, "Warning"); - notification.warning({ + getNotification().warning({ ...COMMON_NOTIFICATION_PROPS, ...cfg, placement: cfg.placement ?? defaultPlacement(), @@ -271,7 +281,7 @@ const NotificationManager = { info(input: string | NotificationConfig) { const cfg = normalize(input, "Info"); - notification.info({ + getNotification().info({ ...COMMON_NOTIFICATION_PROPS, ...cfg, placement: cfg.placement ?? defaultPlacement(), @@ -281,7 +291,7 @@ const NotificationManager = { success(input: string | React.ReactNode | NotificationConfig) { if (React.isValidElement(input)) { - notification.success({ + getNotification().success({ ...COMMON_NOTIFICATION_PROPS, message: "Success", description: input, @@ -291,7 +301,7 @@ const NotificationManager = { return; } const cfg = normalize(input as string | NotificationConfig, "Success"); - notification.success({ + getNotification().success({ ...COMMON_NOTIFICATION_PROPS, ...cfg, placement: cfg.placement ?? defaultPlacement(), @@ -316,11 +326,11 @@ const NotificationManager = { title === "Content Blocked" || title === "Integration Error" ) { - notification.warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 7 }); + getNotification().warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 7 }); return; } if (title === "Server Error") { - notification.error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 8 }); + getNotification().error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 8 }); return; } if ( @@ -331,10 +341,10 @@ const NotificationManager = { title === "Error" || title === "Already Exists" ) { - notification.error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 }); + getNotification().error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 }); return; } - notification.info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 }); + getNotification().info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 }); return; } @@ -343,18 +353,18 @@ const NotificationManager = { const payload = { ...base, message: cls?.title ?? "Info" }; if (cls?.kind === "success") { - notification.success({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 3.5 }); + getNotification().success({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 3.5 }); return; } if (cls?.kind === "warning") { - notification.warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 }); + getNotification().warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 }); return; } - notification.info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 }); + getNotification().info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 }); }, clear() { - notification.destroy(); + getNotification().destroy(); }, }; diff --git a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx new file mode 100644 index 00000000000..b10aa715ed5 --- /dev/null +++ b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx @@ -0,0 +1,25 @@ +"use client"; + +import React, { useEffect } from "react"; +import { App } from "antd"; +import { setNotificationInstance } from "@/components/molecules/notifications_manager"; + +// Inner component to use the hook +const AntdAppInit = () => { + const { notification } = App.useApp(); + + useEffect(() => { + setNotificationInstance(notification); + }, [notification]); + + return null; +}; + +export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) { + return ( + + + {children} + + ); +} From eee37c569ddad6ca13d9055980506e605a77430e Mon Sep 17 00:00:00 2001 From: Swayambhu Date: Fri, 6 Feb 2026 10:01:44 +0530 Subject: [PATCH 04/13] fix: ensure Ant Design notification instance is initialized only once using `useRef`. --- ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx index b10aa715ed5..c5adbff86b2 100644 --- a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx +++ b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx @@ -1,15 +1,19 @@ "use client"; -import React, { useEffect } from "react"; +import React, { useEffect, useRef } from "react"; import { App } from "antd"; import { setNotificationInstance } from "@/components/molecules/notifications_manager"; // Inner component to use the hook const AntdAppInit = () => { const { notification } = App.useApp(); + const initialized = useRef(false); useEffect(() => { - setNotificationInstance(notification); + if (!initialized.current) { + setNotificationInstance(notification); + initialized.current = true; + } }, [notification]); return null; From a48a8ec945d8a5cde7419a22eca3b90493b3d18e Mon Sep 17 00:00:00 2001 From: Swayambhu Date: Fri, 6 Feb 2026 15:35:11 +0530 Subject: [PATCH 05/13] refactor: Directly use Ant Design's `notification` hook instead of `App.useApp` for notification management. --- .../src/contexts/AntdGlobalProvider.tsx | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx index c5adbff86b2..5b5c1036fa9 100644 --- a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx +++ b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx @@ -1,29 +1,24 @@ "use client"; import React, { useEffect, useRef } from "react"; -import { App } from "antd"; +import { notification } from "antd"; import { setNotificationInstance } from "@/components/molecules/notifications_manager"; -// Inner component to use the hook -const AntdAppInit = () => { - const { notification } = App.useApp(); +export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) { + const [api, contextHolder] = notification.useNotification(); const initialized = useRef(false); useEffect(() => { if (!initialized.current) { - setNotificationInstance(notification); + setNotificationInstance(api); initialized.current = true; } - }, [notification]); + }, [api]); - return null; -}; - -export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) { return ( - - + <> + {contextHolder} {children} - + ); } From 6a213fc3bca18c023b535254d3111017f560a3a7 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Sat, 7 Feb 2026 00:00:35 +0100 Subject: [PATCH 06/13] fix(mcp): resolve OAuth2 'Capabilities: none' bug for upstream MCP servers (#20602) - process_mcp_request() now falls back to OAuth2 passthrough when Authorization header contains a non-LiteLLM token (catches HTTPException and ProxyException 401/403) - MCPClient._get_auth_headers() adds missing MCPAuth.oauth2 case --- litellm/experimental_mcp_client/client.py | 2 + .../mcp_server/auth/user_api_key_auth_mcp.py | 52 +++- tests/mcp_tests/test_mcp_client_unit.py | 21 ++ .../auth/test_user_api_key_auth_mcp.py | 229 ++++++++++++++++++ 4 files changed, 293 insertions(+), 11 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3e8f9bc337b..5e21ff9754f 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -209,6 +209,8 @@ class MCPClient: headers["X-API-Key"] = self._mcp_auth_value elif self.auth_type == MCPAuth.authorization: headers["Authorization"] = self._mcp_auth_value + elif self.auth_type == MCPAuth.oauth2: + headers["Authorization"] = f"Bearer {self._mcp_auth_value}" elif isinstance(self._mcp_auth_value, dict): headers.update(self._mcp_auth_value) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7e70b5baae4..786cfbfb008 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,11 +1,12 @@ from typing import Dict, List, Optional, Set, Tuple +from fastapi import HTTPException from starlette.datastructures import Headers from starlette.requests import Request from starlette.types import Scope from litellm._logging import verbose_logger -from litellm.proxy._types import LiteLLM_TeamTable, SpecialHeaders, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_TeamTable, ProxyException, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -63,6 +64,13 @@ class MCPRequestHandler: HTTPException: If headers are invalid or missing required headers """ headers = MCPRequestHandler._safe_get_headers_from_scope(scope) + + # Check if there is an explicit LiteLLM API key (primary header) + has_explicit_litellm_key = ( + headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) + is not None + ) + litellm_api_key = ( MCPRequestHandler.get_litellm_api_key_from_headers(headers) or "" ) @@ -106,16 +114,38 @@ class MCPRequestHandler: request.body = mock_body # type: ignore if ".well-known" in str(request.url): # public routes validated_user_api_key_auth = UserAPIKeyAuth() - # elif litellm_api_key == "": - # from fastapi import HTTPException - - # raise HTTPException( - # status_code=401, - # detail="LiteLLM API key is missing. Please add it or use OAuth authentication.", - # headers={ - # "WWW-Authenticate": f'Bearer resource_metadata=f"{request.base_url}/.well-known/oauth-protected-resource"', - # }, - # ) + elif has_explicit_litellm_key: + # Explicit x-litellm-api-key provided - always validate normally + validated_user_api_key_auth = await user_api_key_auth( + api_key=litellm_api_key, request=request + ) + elif oauth2_headers: + # No x-litellm-api-key, but Authorization header present. + # Could be a LiteLLM key (backward compat) OR an OAuth2 token + # from an upstream MCP provider (e.g. Atlassian). + # Try LiteLLM auth first; on auth failure, treat as OAuth2 passthrough. + try: + validated_user_api_key_auth = await user_api_key_auth( + api_key=litellm_api_key, request=request + ) + except HTTPException as e: + if e.status_code in (401, 403): + verbose_logger.debug( + "MCP OAuth2: Authorization header is not a valid LiteLLM key, " + "treating as OAuth2 token passthrough" + ) + validated_user_api_key_auth = UserAPIKeyAuth() + else: + raise + except ProxyException as e: + if str(e.code) in ("401", "403"): + verbose_logger.debug( + "MCP OAuth2: Authorization header is not a valid LiteLLM key, " + "treating as OAuth2 token passthrough" + ) + validated_user_api_key_auth = UserAPIKeyAuth() + else: + raise else: validated_user_api_key_auth = await user_api_key_auth( api_key=litellm_api_key, request=request diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index 9533e56bcc8..c70d0c42cd8 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -77,6 +77,27 @@ class TestMCPClientUnitTests: "Authorization": "Token custom_token", } + # OAuth2 + client = MCPClient( + "http://example.com", + auth_type=MCPAuth.oauth2, + auth_value="oauth2-access-token-xyz", + ) + headers = client._get_auth_headers() + assert headers == { + "Authorization": "Bearer oauth2-access-token-xyz", + } + + # OAuth2 with extra_headers (per-user flow overrides auth_value) + client = MCPClient( + "http://example.com", + auth_type=MCPAuth.oauth2, + auth_value="static-server-token", + extra_headers={"Authorization": "Bearer per-user-token"}, + ) + headers = client._get_auth_headers() + assert headers["Authorization"] == "Bearer per-user-token" + # No auth client = MCPClient("http://example.com") headers = client._get_auth_headers() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e1e4b3a8b6d..68afe784988 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -544,6 +544,235 @@ class TestMCPRequestHandler: assert mcp_server_auth_headers == {} +@pytest.mark.asyncio +class TestMCPOAuth2AuthFlow: + """Test suite for OAuth2 authentication flow in MCP requests. + + Tests the fix for the 'Capabilities: none' bug where OAuth2 tokens + from upstream MCP providers (e.g., Atlassian) were mistakenly validated + as LiteLLM API keys, causing auth failures and empty tool listings. + """ + + async def test_oauth2_token_in_authorization_header_fallback(self): + """ + When only Authorization header is present with a non-LiteLLM OAuth2 token, + auth should fall back to permissive mode (OAuth2 passthrough). + """ + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/atlassian_mcp", + "headers": [ + (b"authorization", b"Bearer atlassian-oauth2-access-token-xyz"), + ], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ): + ( + auth_result, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = await MCPRequestHandler.process_mcp_request(scope) + + # Should succeed with default UserAPIKeyAuth (OAuth2 fallback) + assert auth_result is not None + assert isinstance(auth_result, UserAPIKeyAuth) + # OAuth2 headers should contain the token for upstream forwarding + assert ( + oauth2_headers.get("Authorization") + == "Bearer atlassian-oauth2-access-token-xyz" + ) + + async def test_explicit_litellm_key_with_oauth2_authorization(self): + """ + When both x-litellm-api-key AND Authorization header are present, + LiteLLM key should be used for auth and Authorization preserved for OAuth2. + """ + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/atlassian_mcp", + "headers": [ + (b"x-litellm-api-key", b"sk-litellm-valid-key"), + (b"authorization", b"Bearer atlassian-oauth2-token"), + ], + } + + async def mock_user_api_key_auth(api_key, request): + return UserAPIKeyAuth(api_key=api_key, user_id="test-user") + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth, + ) as mock_auth: + ( + auth_result, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = await MCPRequestHandler.process_mcp_request(scope) + + # LiteLLM key should be used for auth + mock_auth.assert_called_once() + call_args = mock_auth.call_args + assert call_args.kwargs["api_key"] == "sk-litellm-valid-key" + + # OAuth2 headers should still contain the Authorization token + assert ( + oauth2_headers.get("Authorization") + == "Bearer atlassian-oauth2-token" + ) + + async def test_litellm_key_in_authorization_backward_compat(self): + """ + Backward compatibility: when only Authorization header is present + with a valid LiteLLM key (not OAuth2), auth should succeed normally. + """ + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/some_server", + "headers": [ + (b"authorization", b"Bearer sk-litellm-valid-key"), + ], + } + + async def mock_user_api_key_auth(api_key, request): + return UserAPIKeyAuth(api_key=api_key, user_id="test-user") + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth, + ) as mock_auth: + ( + auth_result, + _, + _, + _, + _, + _, + ) = await MCPRequestHandler.process_mcp_request(scope) + + # Should succeed with the LiteLLM key from Authorization header + assert auth_result.api_key == "Bearer sk-litellm-valid-key" + mock_auth.assert_called_once() + + async def test_non_auth_http_exception_still_raises(self): + """ + If user_api_key_auth raises a non-401/403 HTTPException (e.g., 500), + it should NOT be caught by the OAuth2 fallback. + """ + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/some_server", + "headers": [ + (b"authorization", b"Bearer some-token"), + ], + } + + async def mock_user_api_key_auth_server_error(api_key, request): + raise HTTPException(status_code=500, detail="Internal server error") + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_server_error, + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 500 + + async def test_proxy_exception_oauth2_fallback(self): + """ + user_api_key_auth raises ProxyException (not HTTPException) in production. + The OAuth2 fallback must catch ProxyException with code 401/403 too. + """ + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/atlassian_mcp", + "headers": [ + (b"authorization", b"Bearer atlassian-oauth2-access-token-xyz"), + ], + } + + async def mock_user_api_key_auth_proxy_exception(api_key, request): + raise ProxyException( + message="Authentication Error: Invalid API key", + type="auth_error", + param="api_key", + code=401, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_proxy_exception, + ): + ( + auth_result, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = await MCPRequestHandler.process_mcp_request(scope) + + # Should succeed with default UserAPIKeyAuth (OAuth2 fallback) + assert auth_result is not None + assert isinstance(auth_result, UserAPIKeyAuth) + assert ( + oauth2_headers.get("Authorization") + == "Bearer atlassian-oauth2-access-token-xyz" + ) + + async def test_proxy_exception_non_auth_still_raises(self): + """ + ProxyException with non-401/403 code should NOT be caught. + """ + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/some_server", + "headers": [ + (b"authorization", b"Bearer some-token"), + ], + } + + async def mock_user_api_key_auth_500(api_key, request): + raise ProxyException( + message="Internal error", + type="server_error", + param=None, + code=500, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_500, + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request(scope) + + class TestMCPCustomHeaderName: """Test suite for custom MCP authentication header name functionality""" From b78f4c924c601036da005ba1194c05f95f319009 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 6 Feb 2026 15:02:34 -0800 Subject: [PATCH 07/13] [Fix] A2a Agent Gateway Fixes - A2A agents deployed with localhost/internal URLs in their agent cards (e.g., http://0.0.0.0:8001/) (#20604) * v1 card resolver fix * fix: is_localhost_or_internal_url * fix code * test_fix_agent_card_url_replaces_localhost * test restruct * test_a2a_non_streaming * test agnts * add exception handling * init errors * add localhost retry * add agent_testing * test_a2a_non_streaming * _build_streaming_logging_obj * code qa fixes * test_card_resolver_fallback_from_new_to_old_path * fix linting --- .circleci/config.yml | 53 +++++ litellm/a2a_protocol/__init__.py | 14 ++ litellm/a2a_protocol/card_resolver.py | 67 +++++- .../a2a_protocol/exception_mapping_utils.py | 203 ++++++++++++++++++ litellm/a2a_protocol/exceptions.py | 150 +++++++++++++ litellm/a2a_protocol/main.py | 184 ++++++++++++---- litellm/constants.py | 16 ++ .../local_vertex_agent.py | 0 .../{ => local_only_agent_tests}/test_a2a.py | 0 .../test_a2a_completion_bridge.py | 0 tests/agent_tests/test_a2a_agent.py | 83 +++++++ .../a2a_protocol/test_card_resolver.py | 85 +++++--- 12 files changed, 773 insertions(+), 82 deletions(-) create mode 100644 litellm/a2a_protocol/exception_mapping_utils.py create mode 100644 litellm/a2a_protocol/exceptions.py rename tests/agent_tests/{ => local_only_agent_tests}/local_vertex_agent.py (100%) rename tests/agent_tests/{ => local_only_agent_tests}/test_a2a.py (100%) rename tests/agent_tests/{ => local_only_agent_tests}/test_a2a_completion_bridge.py (100%) create mode 100644 tests/agent_tests/test_a2a_agent.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 8672561f654..e171759f1c4 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1372,6 +1372,51 @@ jobs: paths: - mcp_coverage.xml - mcp_coverage + agent_testing: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + + steps: + - checkout + - setup_google_dns + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-cov==5.0.0" + pip install "pytest-asyncio==0.21.1" + pip install "respx==0.22.0" + pip install "pydantic==2.11.0" + pip install "a2a-sdk" + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/agent_tests --ignore=tests/agent_tests/local_only_agent_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml agent_coverage.xml + mv .coverage agent_coverage + + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - agent_coverage.xml + - agent_coverage guardrails_testing: docker: - image: cimg/python:3.11 @@ -4264,6 +4309,12 @@ workflows: only: - main - /litellm_.*/ + - agent_testing: + filters: + branches: + only: + - main + - /litellm_.*/ - guardrails_testing: filters: branches: @@ -4371,6 +4422,7 @@ workflows: - llm_translation_testing - realtime_translation_testing - mcp_testing + - agent_testing - google_generate_content_endpoint_testing - guardrails_testing - llm_responses_api_testing @@ -4449,6 +4501,7 @@ workflows: - llm_translation_testing - realtime_translation_testing - mcp_testing + - agent_testing - google_generate_content_endpoint_testing - llm_responses_api_testing - ocr_testing diff --git a/litellm/a2a_protocol/__init__.py b/litellm/a2a_protocol/__init__.py index d8d349bb98a..85c03687e25 100644 --- a/litellm/a2a_protocol/__init__.py +++ b/litellm/a2a_protocol/__init__.py @@ -39,6 +39,12 @@ Example usage (class-based): """ from litellm.a2a_protocol.client import A2AClient +from litellm.a2a_protocol.exceptions import ( + A2AAgentCardError, + A2AConnectionError, + A2AError, + A2ALocalhostURLError, +) from litellm.a2a_protocol.main import ( aget_agent_card, asend_message, @@ -49,11 +55,19 @@ from litellm.a2a_protocol.main import ( from litellm.types.agents import LiteLLMSendMessageResponse __all__ = [ + # Client "A2AClient", + # Functions "asend_message", "send_message", "asend_message_streaming", "aget_agent_card", "create_a2a_client", + # Response types "LiteLLMSendMessageResponse", + # Exceptions + "A2AError", + "A2AConnectionError", + "A2AAgentCardError", + "A2ALocalhostURLError", ] diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 7c4c5af149d..4c5dd3e3ba6 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -7,6 +7,7 @@ Extends the A2A SDK's card resolver to support multiple well-known paths. from typing import TYPE_CHECKING, Any, Dict, Optional from litellm._logging import verbose_logger +from litellm.constants import LOCALHOST_URL_PATTERNS if TYPE_CHECKING: from a2a.types import AgentCard @@ -26,15 +27,61 @@ except ImportError: pass +def is_localhost_or_internal_url(url: Optional[str]) -> bool: + """ + Check if a URL is a localhost or internal URL. + + This detects common development URLs that are accidentally left in + agent cards when deploying to production. + + Args: + url: The URL to check + + Returns: + True if the URL is localhost/internal + """ + if not url: + return False + + url_lower = url.lower() + + return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS) + + +def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard": + """ + Fix the agent card URL if it contains a localhost/internal address. + + Many A2A agents are deployed with agent cards that contain internal URLs + like "http://0.0.0.0:8001/" or "http://localhost:8000/". This function + replaces such URLs with the provided base_url. + + Args: + agent_card: The agent card to fix + base_url: The base URL to use as replacement + + Returns: + The agent card with the URL fixed if necessary + """ + card_url = getattr(agent_card, "url", None) + + if card_url and is_localhost_or_internal_url(card_url): + # Normalize base_url to ensure it ends with / + fixed_url = base_url.rstrip("/") + "/" + agent_card.url = fixed_url + + return agent_card + + class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] """ Custom A2A card resolver that supports multiple well-known paths. - + Extends the base A2ACardResolver to try both: - /.well-known/agent-card.json (standard) - /.well-known/agent.json (previous/alternative) """ - + async def get_agent_card( self, relative_card_path: Optional[str] = None, @@ -42,17 +89,17 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] ) -> "AgentCard": """ Fetch the agent card, trying multiple well-known paths. - + First tries the standard path, then falls back to the previous path. - + Args: relative_card_path: Optional path to the agent card endpoint. If None, tries both well-known paths. http_kwargs: Optional dictionary of keyword arguments to pass to httpx.get - + Returns: AgentCard from the A2A agent - + Raises: A2AClientHTTPError or A2AClientJSONError if both paths fail """ @@ -62,13 +109,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] relative_card_path=relative_card_path, http_kwargs=http_kwargs, ) - + # Try both well-known paths paths = [ AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH, ] - + last_error = None for path in paths: try: @@ -85,11 +132,11 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] ) last_error = e continue - + # If we get here, all paths failed - re-raise the last error if last_error is not None: raise last_error - + # This shouldn't happen, but just in case raise Exception( f"Failed to fetch agent card from {self.base_url}. " diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py new file mode 100644 index 00000000000..8463080b358 --- /dev/null +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -0,0 +1,203 @@ +""" +A2A Protocol Exception Mapping Utils. + +Maps A2A SDK exceptions to LiteLLM A2A exception types. +""" + +from typing import TYPE_CHECKING, Any, Optional + +from litellm._logging import verbose_logger +from litellm.a2a_protocol.card_resolver import ( + fix_agent_card_url, + is_localhost_or_internal_url, +) +from litellm.a2a_protocol.exceptions import ( + A2AAgentCardError, + A2AConnectionError, + A2AError, + A2ALocalhostURLError, +) +from litellm.constants import CONNECTION_ERROR_PATTERNS + +if TYPE_CHECKING: + from a2a.client import A2AClient as A2AClientType + + +# Runtime import +A2A_SDK_AVAILABLE = False +try: + from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef] + + A2A_SDK_AVAILABLE = True +except ImportError: + _A2AClient = None # type: ignore[misc] + + +class A2AExceptionCheckers: + """ + Helper class for checking various A2A error conditions. + """ + + @staticmethod + def is_connection_error(error_str: str) -> bool: + """ + Check if an error string indicates a connection error. + + Args: + error_str: The error string to check + + Returns: + True if the error indicates a connection issue + """ + if not isinstance(error_str, str): + return False + + error_str_lower = error_str.lower() + return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS) + + @staticmethod + def is_localhost_url(url: Optional[str]) -> bool: + """ + Check if a URL is a localhost/internal URL. + + Args: + url: The URL to check + + Returns: + True if the URL is localhost/internal + """ + return is_localhost_or_internal_url(url) + + @staticmethod + def is_agent_card_error(error_str: str) -> bool: + """ + Check if an error string indicates an agent card error. + + Args: + error_str: The error string to check + + Returns: + True if the error is related to agent card fetching/parsing + """ + if not isinstance(error_str, str): + return False + + error_str_lower = error_str.lower() + agent_card_patterns = [ + "agent card", + "agent-card", + ".well-known", + "card not found", + "invalid agent", + ] + return any(pattern in error_str_lower for pattern in agent_card_patterns) + + +def map_a2a_exception( + original_exception: Exception, + card_url: Optional[str] = None, + api_base: Optional[str] = None, + model: Optional[str] = None, +) -> Exception: + """ + Map an A2A SDK exception to a LiteLLM A2A exception type. + + Args: + original_exception: The original exception from the A2A SDK + card_url: The URL from the agent card (if available) + api_base: The original API base URL + model: The model/agent name + + Returns: + A mapped LiteLLM A2A exception + + Raises: + A2ALocalhostURLError: If the error is a connection error to a localhost URL + A2AConnectionError: If the error is a general connection error + A2AAgentCardError: If the error is related to agent card issues + A2AError: For other A2A-related errors + """ + error_str = str(original_exception) + + # Check for localhost URL connection error (special case - retryable) + if ( + card_url + and api_base + and A2AExceptionCheckers.is_localhost_url(card_url) + and A2AExceptionCheckers.is_connection_error(error_str) + ): + raise A2ALocalhostURLError( + localhost_url=card_url, + base_url=api_base, + original_error=original_exception, + model=model, + ) + + # Check for agent card errors + if A2AExceptionCheckers.is_agent_card_error(error_str): + raise A2AAgentCardError( + message=error_str, + url=api_base, + model=model, + ) + + # Check for general connection errors + if A2AExceptionCheckers.is_connection_error(error_str): + raise A2AConnectionError( + message=error_str, + url=card_url or api_base, + model=model, + ) + + # Default: wrap in generic A2AError + raise A2AError( + message=error_str, + model=model, + ) + + +def handle_a2a_localhost_retry( + error: A2ALocalhostURLError, + agent_card: Any, + a2a_client: "A2AClientType", + is_streaming: bool = False, +) -> "A2AClientType": + """ + Handle A2ALocalhostURLError by fixing the URL and creating a new client. + + This is called when we catch an A2ALocalhostURLError and want to retry + with the corrected URL. + + Args: + error: The localhost URL error + agent_card: The agent card object to fix + a2a_client: The current A2A client + is_streaming: Whether this is a streaming request (for logging) + + Returns: + A new A2A client with the fixed URL + + Raises: + ImportError: If the A2A SDK is not installed + """ + if not A2A_SDK_AVAILABLE or _A2AClient is None: + raise ImportError( + "A2A SDK is required for localhost retry handling. " + "Install it with: pip install a2a" + ) + + request_type = "streaming " if is_streaming else "" + verbose_logger.warning( + f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. " + f"Agent card contains localhost/internal URL. " + f"Retrying with base_url '{error.base_url}'." + ) + + # Fix the agent card URL + fix_agent_card_url(agent_card, error.base_url) + + # Create a new client with the fixed agent card (transport caches URL) + return _A2AClient( + httpx_client=a2a_client._transport.httpx_client, # type: ignore[union-attr] + agent_card=agent_card, + ) diff --git a/litellm/a2a_protocol/exceptions.py b/litellm/a2a_protocol/exceptions.py new file mode 100644 index 00000000000..546b23105be --- /dev/null +++ b/litellm/a2a_protocol/exceptions.py @@ -0,0 +1,150 @@ +""" +A2A Protocol Exceptions. + +Custom exception types for A2A protocol operations, following LiteLLM's exception pattern. +""" + +from typing import Optional + +import httpx + + +class A2AError(Exception): + """ + Base exception for A2A protocol errors. + + Follows the same pattern as LiteLLM's main exceptions. + """ + + def __init__( + self, + message: str, + status_code: int = 500, + llm_provider: str = "a2a_agent", + model: Optional[str] = None, + response: Optional[httpx.Response] = None, + litellm_debug_info: Optional[str] = None, + max_retries: Optional[int] = None, + num_retries: Optional[int] = None, + ): + self.status_code = status_code + self.message = f"litellm.A2AError: {message}" + self.llm_provider = llm_provider + self.model = model + self.litellm_debug_info = litellm_debug_info + self.max_retries = max_retries + self.num_retries = num_retries + self.response = response or httpx.Response( + status_code=self.status_code, + request=httpx.Request(method="POST", url="https://litellm.ai"), + ) + super().__init__(self.message) + + def __str__(self) -> str: + _message = self.message + if self.num_retries: + _message += f" LiteLLM Retried: {self.num_retries} times" + if self.max_retries: + _message += f", LiteLLM Max Retries: {self.max_retries}" + return _message + + def __repr__(self) -> str: + return self.__str__() + + +class A2AConnectionError(A2AError): + """ + Raised when connection to an A2A agent fails. + + This typically occurs when: + - The agent is unreachable + - The agent card contains a localhost/internal URL + - Network issues prevent connection + """ + + def __init__( + self, + message: str, + url: Optional[str] = None, + model: Optional[str] = None, + response: Optional[httpx.Response] = None, + litellm_debug_info: Optional[str] = None, + max_retries: Optional[int] = None, + num_retries: Optional[int] = None, + ): + self.url = url + super().__init__( + message=message, + status_code=503, + llm_provider="a2a_agent", + model=model, + response=response, + litellm_debug_info=litellm_debug_info, + max_retries=max_retries, + num_retries=num_retries, + ) + + +class A2AAgentCardError(A2AError): + """ + Raised when there's an issue with the agent card. + + This includes: + - Failed to fetch agent card + - Invalid agent card format + - Missing required fields + """ + + def __init__( + self, + message: str, + url: Optional[str] = None, + model: Optional[str] = None, + response: Optional[httpx.Response] = None, + litellm_debug_info: Optional[str] = None, + ): + self.url = url + super().__init__( + message=message, + status_code=404, + llm_provider="a2a_agent", + model=model, + response=response, + litellm_debug_info=litellm_debug_info, + ) + + +class A2ALocalhostURLError(A2AConnectionError): + """ + Raised when an agent card contains a localhost/internal URL. + + Many A2A agents are deployed with agent cards that contain internal URLs + like "http://0.0.0.0:8001/" or "http://localhost:8000/". This error + indicates that the URL needs to be corrected and the request should be retried. + + Attributes: + localhost_url: The localhost/internal URL found in the agent card + base_url: The public base URL that should be used instead + original_error: The original connection error that was raised + """ + + def __init__( + self, + localhost_url: str, + base_url: str, + original_error: Optional[Exception] = None, + model: Optional[str] = None, + ): + self.localhost_url = localhost_url + self.base_url = base_url + self.original_error = original_error + + message = ( + f"Agent card contains localhost/internal URL '{localhost_url}'. " + f"Retrying with base URL '{base_url}'." + ) + super().__init__( + message=message, + url=localhost_url, + model=model, + ) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index b326f9e7ed5..642dfaf023c 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -44,6 +44,11 @@ except ImportError: # Import our custom card resolver that supports multiple well-known paths from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver +from litellm.a2a_protocol.exception_mapping_utils import ( + handle_a2a_localhost_retry, + map_a2a_exception, +) +from litellm.a2a_protocol.exceptions import A2ALocalhostURLError # Use our custom resolver instead of the default A2A SDK resolver A2ACardResolver = LiteLLMA2ACardResolver @@ -244,10 +249,50 @@ async def asend_message( verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}") - a2a_response = await a2a_client.send_message(request) + # Get agent card URL for localhost retry logic + agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr( + a2a_client, "agent_card", None + ) + card_url = getattr(agent_card, "url", None) if agent_card else None + + # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL + a2a_response = None + for _ in range(2): # max 2 attempts: original + 1 retry + try: + a2a_response = await a2a_client.send_message(request) + break # success, exit retry loop + except A2ALocalhostURLError as e: + # Localhost URL error - fix and retry + a2a_client = handle_a2a_localhost_retry( + error=e, + agent_card=agent_card, + a2a_client=a2a_client, + is_streaming=False, + ) + card_url = agent_card.url if agent_card else None + except Exception as e: + # Map exception - will raise A2ALocalhostURLError if applicable + try: + map_a2a_exception(e, card_url, api_base, model=agent_name) + except A2ALocalhostURLError as localhost_err: + # Localhost URL error - fix and retry + a2a_client = handle_a2a_localhost_retry( + error=localhost_err, + agent_card=agent_card, + a2a_client=a2a_client, + is_streaming=False, + ) + card_url = agent_card.url if agent_card else None + continue + except Exception: + # Re-raise the mapped exception + raise verbose_logger.info(f"A2A send_message completed, request_id={request.id}") + # a2a_response is guaranteed to be set if we reach here (loop breaks on success or raises) + assert a2a_response is not None + # Wrap in LiteLLM response type for _hidden_params support response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response) @@ -307,6 +352,48 @@ def send_message( ) +def _build_streaming_logging_obj( + request: "SendStreamingMessageRequest", + agent_name: str, + agent_id: Optional[str], + litellm_params: Optional[Dict[str, Any]], + metadata: Optional[Dict[str, Any]], + proxy_server_request: Optional[Dict[str, Any]], +) -> Logging: + """Build logging object for streaming A2A requests.""" + start_time = datetime.datetime.now() + model = f"a2a_agent/{agent_name}" + + logging_obj = Logging( + model=model, + messages=[{"role": "user", "content": "streaming-request"}], + stream=False, + call_type="asend_message_streaming", + start_time=start_time, + litellm_call_id=str(request.id), + function_id=str(request.id), + ) + logging_obj.model = model + logging_obj.custom_llm_provider = "a2a_agent" + logging_obj.model_call_details["model"] = model + logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent" + if agent_id: + logging_obj.model_call_details["agent_id"] = agent_id + + _litellm_params = litellm_params.copy() if litellm_params else {} + if metadata: + _litellm_params["metadata"] = metadata + if proxy_server_request: + _litellm_params["proxy_server_request"] = proxy_server_request + + logging_obj.litellm_params = _litellm_params + logging_obj.optional_params = _litellm_params + logging_obj.model_call_details["litellm_params"] = _litellm_params + logging_obj.model_call_details["metadata"] = metadata or {} + + return logging_obj + + async def asend_message_streaming( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendStreamingMessageRequest"] = None, @@ -403,55 +490,72 @@ async def asend_message_streaming( verbose_logger.info(f"A2A send_message_streaming request_id={request.id}") - # Track for logging - start_time = datetime.datetime.now() - stream = a2a_client.send_message_streaming(request) - # Build logging object for streaming completion callbacks agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr( a2a_client, "agent_card", None ) + card_url = getattr(agent_card, "url", None) if agent_card else None agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown" - model = f"a2a_agent/{agent_name}" - logging_obj = Logging( - model=model, - messages=[{"role": "user", "content": "streaming-request"}], - stream=False, # complete response logging after stream ends - call_type="asend_message_streaming", - start_time=start_time, - litellm_call_id=str(request.id), - function_id=str(request.id), - ) - logging_obj.model = model - logging_obj.custom_llm_provider = "a2a_agent" - logging_obj.model_call_details["model"] = model - logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent" - if agent_id: - logging_obj.model_call_details["agent_id"] = agent_id - - # Propagate litellm_params for spend logging (includes cost_per_query, etc.) - _litellm_params = litellm_params.copy() if litellm_params else {} - # Merge metadata into litellm_params.metadata (required for proxy cost tracking) - if metadata: - _litellm_params["metadata"] = metadata - if proxy_server_request: - _litellm_params["proxy_server_request"] = proxy_server_request - - logging_obj.litellm_params = _litellm_params - logging_obj.optional_params = _litellm_params # used by cost calc - logging_obj.model_call_details["litellm_params"] = _litellm_params - logging_obj.model_call_details["metadata"] = metadata or {} - - iterator = A2AStreamingIterator( - stream=stream, + logging_obj = _build_streaming_logging_obj( request=request, - logging_obj=logging_obj, agent_name=agent_name, + agent_id=agent_id, + litellm_params=litellm_params, + metadata=metadata, + proxy_server_request=proxy_server_request, ) - async for chunk in iterator: - yield chunk + # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL + # Connection errors in streaming typically occur on first chunk iteration + first_chunk = True + for attempt in range(2): # max 2 attempts: original + 1 retry + stream = a2a_client.send_message_streaming(request) + iterator = A2AStreamingIterator( + stream=stream, + request=request, + logging_obj=logging_obj, + agent_name=agent_name, + ) + + try: + first_chunk = True + async for chunk in iterator: + if first_chunk: + first_chunk = False # connection succeeded + yield chunk + return # stream completed successfully + except A2ALocalhostURLError as e: + # Only retry on first chunk, not mid-stream + if first_chunk and attempt == 0: + a2a_client = handle_a2a_localhost_retry( + error=e, + agent_card=agent_card, + a2a_client=a2a_client, + is_streaming=True, + ) + card_url = agent_card.url if agent_card else None + else: + raise + except Exception as e: + # Only map exception on first chunk + if first_chunk and attempt == 0: + try: + map_a2a_exception(e, card_url, api_base, model=agent_name) + except A2ALocalhostURLError as localhost_err: + # Localhost URL error - fix and retry + a2a_client = handle_a2a_localhost_retry( + error=localhost_err, + agent_card=agent_card, + a2a_client=a2a_client, + is_streaming=True, + ) + card_url = agent_card.url if agent_card else None + continue + except Exception: + # Re-raise the mapped exception + raise + raise async def create_a2a_client( diff --git a/litellm/constants.py b/litellm/constants.py index 25decd363a6..3c618723d64 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -306,6 +306,22 @@ DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2 #### Networking settings #### request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes +# Patterns that indicate a localhost/internal URL in A2A agent cards that should be +# replaced with the original base_url. This is a common misconfiguration where +# developers deploy agents with development URLs in their agent cards. +LOCALHOST_URL_PATTERNS: List[str] = [ + "localhost", + "127.0.0.1", + "0.0.0.0", + "[::1]", # IPv6 localhost +] +# Patterns in error messages that indicate a connection failure +CONNECTION_ERROR_PATTERNS: List[str] = [ + "connect", + "connection", + "network", + "refused", +] STREAM_SSE_DONE_STRING: str = "[DONE]" STREAM_SSE_DATA_PREFIX: str = "data: " ### SPEND TRACKING ### diff --git a/tests/agent_tests/local_vertex_agent.py b/tests/agent_tests/local_only_agent_tests/local_vertex_agent.py similarity index 100% rename from tests/agent_tests/local_vertex_agent.py rename to tests/agent_tests/local_only_agent_tests/local_vertex_agent.py diff --git a/tests/agent_tests/test_a2a.py b/tests/agent_tests/local_only_agent_tests/test_a2a.py similarity index 100% rename from tests/agent_tests/test_a2a.py rename to tests/agent_tests/local_only_agent_tests/test_a2a.py diff --git a/tests/agent_tests/test_a2a_completion_bridge.py b/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py similarity index 100% rename from tests/agent_tests/test_a2a_completion_bridge.py rename to tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py new file mode 100644 index 00000000000..acaf615ec88 --- /dev/null +++ b/tests/agent_tests/test_a2a_agent.py @@ -0,0 +1,83 @@ +""" +Simple A2A agent tests - non-streaming and streaming. + +These tests validate the localhost URL retry logic: if an A2A agent's card +contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request +will fail with a connection error. LiteLLM detects this and automatically +retries using the original api_base URL instead. + +Requires A2A_AGENT_URL environment variable to be set. + +Run with: + A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s +""" + +import os + +import pytest +from uuid import uuid4 + + +def get_a2a_agent_url(): + """Get A2A agent URL from environment, skip test if not set.""" + url = os.environ.get("A2A_AGENT_URL") + return url + + +@pytest.mark.asyncio +async def test_a2a_non_streaming(): + """Test non-streaming A2A request.""" + from a2a.types import MessageSendParams, SendMessageRequest + from litellm.a2a_protocol import asend_message + + api_base = get_a2a_agent_url() + + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Say hello in one word"}], + "messageId": uuid4().hex, + } + ), + ) + + response = await asend_message( + request=request, + api_base=api_base, + ) + + assert response is not None + print(f"\nNon-streaming response: {response}") + + +@pytest.mark.asyncio +async def test_a2a_streaming(): + """Test streaming A2A request.""" + from a2a.types import MessageSendParams, SendStreamingMessageRequest + from litellm.a2a_protocol import asend_message_streaming + + api_base = get_a2a_agent_url() + + request = SendStreamingMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Say hello in one word"}], + "messageId": uuid4().hex, + } + ), + ) + + chunks = [] + async for chunk in asend_message_streaming( + request=request, + api_base=api_base, + ): + chunks.append(chunk) + print(f"\nStreaming chunk: {chunk}") + + assert len(chunks) > 0, "Should receive at least one chunk" + print(f"\nTotal chunks received: {len(chunks)}") diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/test_litellm/a2a_protocol/test_card_resolver.py index a1bd33107b6..1bdab50860c 100644 --- a/tests/test_litellm/a2a_protocol/test_card_resolver.py +++ b/tests/test_litellm/a2a_protocol/test_card_resolver.py @@ -4,11 +4,16 @@ Mock tests for LiteLLMA2ACardResolver. Tests that the card resolver tries both old and new well-known paths. """ -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.a2a_protocol.card_resolver import ( + LiteLLMA2ACardResolver, + fix_agent_card_url, + is_localhost_or_internal_url, +) + @pytest.mark.asyncio async def test_card_resolver_fallback_from_new_to_old_path(): @@ -24,39 +29,31 @@ async def test_card_resolver_fallback_from_new_to_old_path(): # Track which paths were called paths_called = [] - # Create a mock base class - class MockA2ACardResolver: - def __init__(self, base_url): - self.base_url = base_url - - async def get_agent_card(self, relative_card_path=None, http_kwargs=None): - paths_called.append(relative_card_path) - if relative_card_path == "/.well-known/agent-card.json": - # First call (new path) fails - raise Exception("404 Not Found") - else: - # Second call (old path) succeeds - return mock_agent_card - - # Create mock A2A module - mock_a2a_module = MagicMock() - mock_a2a_client = MagicMock() - mock_a2a_constants = MagicMock() - mock_a2a_constants.AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json" - mock_a2a_constants.PREV_AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json" - - with patch.dict( - sys.modules, - { - "a2a": mock_a2a_module, - "a2a.client": MagicMock(A2ACardResolver=MockA2ACardResolver), - "a2a.utils.constants": mock_a2a_constants, - }, + # Create a mock for the parent's get_agent_card method + async def mock_parent_get_agent_card( + self, relative_card_path=None, http_kwargs=None ): - # Import after patching - from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver + paths_called.append(relative_card_path) + if relative_card_path == "/.well-known/agent-card.json": + # First call (new path) fails + raise Exception("404 Not Found") + else: + # Second call (old path) succeeds + return mock_agent_card - resolver = LiteLLMA2ACardResolver(base_url="http://test-agent:8000") + # Create a mock httpx client + mock_httpx_client = MagicMock() + + # Patch the parent class's get_agent_card method + # We need to patch the actual parent class method that super() calls + with patch.object( + LiteLLMA2ACardResolver.__bases__[0], + "get_agent_card", + mock_parent_get_agent_card, + ): + resolver = LiteLLMA2ACardResolver( + httpx_client=mock_httpx_client, base_url="http://test-agent:8000" + ) result = await resolver.get_agent_card() # Verify both paths were tried in correct order @@ -67,3 +64,27 @@ async def test_card_resolver_fallback_from_new_to_old_path(): # Verify the result assert result == mock_agent_card assert result.name == "Test Agent" + + +def test_is_localhost_or_internal_url(): + """Test that localhost/internal URLs are correctly detected.""" + # Should return True for localhost variants + assert is_localhost_or_internal_url("http://localhost:8000/") is True + assert is_localhost_or_internal_url("http://0.0.0.0:8001/") is True + + # Should return False for public URLs + assert is_localhost_or_internal_url("https://my-agent.example.com/") is False + assert is_localhost_or_internal_url(None) is False + + +def test_fix_agent_card_url_replaces_localhost(): + """Test that fix_agent_card_url replaces localhost URLs with base_url.""" + # Create mock agent card with localhost URL + mock_card = MagicMock() + mock_card.url = "http://0.0.0.0:8001/" + + # Fix the URL + result = fix_agent_card_url(mock_card, "https://my-public-agent.example.com") + + # Verify localhost URL was replaced with base_url + assert result.url == "https://my-public-agent.example.com/" From fd3ca081cc67e2d641ac8b988aeb37b6baff2a46 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 6 Feb 2026 15:07:29 -0800 Subject: [PATCH 08/13] use cached keys and teams for router settings --- litellm/proxy/_types.py | 1 + litellm/proxy/common_request_processing.py | 8 +- litellm/proxy/proxy_server.py | 80 +++++++++++++++++++ .../proxy/test_common_request_processing.py | 4 +- 4 files changed, 87 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f38f94f4c98..501cfacaf91 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2162,6 +2162,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") last_rotation_at: Optional[datetime] = None # When this key was last rotated key_rotation_at: Optional[datetime] = None # When this key should next be rotated + router_settings: Optional[dict] = None model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ee38d56c855..3e079045774 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -616,19 +616,21 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type # type: ignore ) - # Apply hierarchical router_settings (Key > Team > Global) + # Apply hierarchical router_settings (Key > Team) + # Global router_settings are already on the Router object itself. if llm_router is not None and proxy_config is not None: from litellm.proxy.proxy_server import prisma_client router_settings = await proxy_config._get_hierarchical_router_settings( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, ) - # If router_settings found (from key, team, or global), apply them + # If router_settings found (from key or team), apply them # Pass settings as per-request overrides instead of creating a new Router # This avoids expensive Router instantiation on each request - if router_settings is not None and router_settings: + if router_settings is not None: self.data["router_settings_override"] = router_settings if "messages" in self.data and self.data["messages"]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1a4db938474..60327d8da17 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3403,6 +3403,86 @@ class ProxyConfig: decrypted_variables[k] = decrypted_value return decrypted_variables + @staticmethod + def _parse_router_settings_value(value: Any) -> Optional[dict]: + """ + Parse a router_settings value that may be a dict or a JSON/YAML string. + + Returns a non-empty dict if valid, otherwise None. + """ + if value is None: + return None + + parsed: Optional[dict] = None + if isinstance(value, dict): + parsed = value + elif isinstance(value, str): + import json + + import yaml + + try: + parsed = yaml.safe_load(value) + except (yaml.YAMLError, json.JSONDecodeError): + try: + parsed = json.loads(value) + except json.JSONDecodeError: + pass + + if isinstance(parsed, dict) and parsed: + return parsed + return None + + async def _get_hierarchical_router_settings( + self, + user_api_key_dict: Optional["UserAPIKeyAuth"], + prisma_client: Optional[PrismaClient], + proxy_logging_obj: Optional["ProxyLogging"] = None, + ) -> Optional[dict]: + """ + Get router_settings in priority order: Key > Team + + Uses the already-cached key object and the cached team lookup + (get_team_object) to avoid direct DB queries on the hot path. + + Global router_settings are NOT looked up here — they are already + applied to the Router object at config-load / DB-sync time. + + Returns: + dict: router_settings, or None if no settings found + """ + # 1. Try key-level router_settings + # user_api_key_dict is already the cached/authenticated key object — + # no DB call needed. + if user_api_key_dict is not None: + key_settings = self._parse_router_settings_value( + getattr(user_api_key_dict, "router_settings", None) + ) + if key_settings is not None: + return key_settings + + # 2. Try team-level router_settings using cached team lookup + # get_team_object checks in-memory cache / Redis first, only falls + # back to DB on a cache miss. + if user_api_key_dict is not None and user_api_key_dict.team_id is not None: + try: + team_obj = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + team_settings = self._parse_router_settings_value( + getattr(team_obj, "router_settings", None) + ) + if team_settings is not None: + return team_settings + except Exception: + # If team lookup fails, no team-level settings available + pass + + return None + async def _add_router_settings_from_db_config( self, config_data: dict, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c45ffe0ada0..7bebe00d61e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -78,7 +78,6 @@ class TestProxyBaseLLMRequestProcessing: assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] @pytest.mark.asyncio -<<<<<<< HEAD async def test_should_apply_hierarchical_router_settings_as_override( self, monkeypatch ): @@ -148,6 +147,7 @@ class TestProxyBaseLLMRequestProcessing: mock_proxy_config._get_hierarchical_router_settings.assert_called_once_with( user_api_key_dict=mock_user_api_key_dict, prisma_client=mock_prisma_client, + proxy_logging_obj=mock_proxy_logging_obj, ) # get_model_list should NOT be called - we no longer copy model list for per-request routers mock_llm_router.get_model_list.assert_not_called() @@ -165,8 +165,6 @@ class TestProxyBaseLLMRequestProcessing: assert "model_list" not in router_settings_override @pytest.mark.asyncio -======= ->>>>>>> origin async def test_stream_timeout_header_processing(self): """ Test that x-litellm-stream-timeout header gets processed and added to request data as stream_timeout. From f4a0b80a25e94d3f034afab6e598d7c88be19565 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Sat, 7 Feb 2026 01:05:51 +0100 Subject: [PATCH 09/13] fix(sso): extract user roles from JWT access token for Keycloak compatibility (#20591) Keycloak (and similar OIDC providers) include role claims in the JWT access token but not in the UserInfo endpoint response. Previously, roles were only extracted from UserInfo, causing all SSO users to default to internal_user_view_only regardless of their actual role. Changes: - Extract user roles from JWT access token in process_sso_jwt_access_token() when UserInfo doesn't provide them (tries role_mappings first, then GENERIC_USER_ROLE_ATTRIBUTE) - Handle list-type role values in get_litellm_user_role() since Keycloak returns roles as arrays (e.g. ["proxy_admin"] instead of "proxy_admin") - Add 9 new unit tests covering role extraction and list handling - Update 3 existing tests for new JWT decode behavior Closes #20407 --- litellm/proxy/management_endpoints/types.py | 13 +- litellm/proxy/management_endpoints/ui_sso.py | 84 +++++-- .../proxy/management_endpoints/test_ui_sso.py | 231 +++++++++++++++++- 3 files changed, 304 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index ad2ad0a5fe5..a35fc4a5f3f 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -28,17 +28,24 @@ def is_valid_litellm_user_role(role_str: str) -> bool: return False -def get_litellm_user_role(role_str: str) -> Optional[LitellmUserRoles]: +def get_litellm_user_role(role_str) -> Optional[LitellmUserRoles]: """ - Convert a string to a LitellmUserRoles enum if valid (case-insensitive). + Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive). + + Handles list inputs since some SSO providers (e.g., Keycloak) return roles + as arrays like ["proxy_admin"] instead of plain strings. Args: - role_str: String to convert (e.g., "proxy_admin", "PROXY_ADMIN", "internal_user") + role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"]) Returns: LitellmUserRoles enum if valid, None otherwise """ try: + if isinstance(role_str, list): + if len(role_str) == 0: + return None + role_str = role_str[0] # Use _value2member_map_ for O(1) lookup, case-insensitive result = LitellmUserRoles._value2member_map_.get(role_str.lower()) return cast(Optional[LitellmUserRoles], result) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 2d248dc81f3..278f3bdaafd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -171,36 +171,86 @@ def process_sso_jwt_access_token( access_token_str: Optional[str], sso_jwt_handler: Optional[JWTHandler], result: Union[OpenID, dict, None], + role_mappings: Optional["RoleMappings"] = None, ) -> None: """ - Process SSO JWT access token and extract team IDs if available. + Process SSO JWT access token and extract team IDs and user role if available. - This function decodes the JWT access token and extracts team IDs using the - sso_jwt_handler, then sets the team_ids attribute on the result object. + This function decodes the JWT access token and extracts team IDs and user + role, then sets them on the result object. Role extraction from the access + token is needed because some SSO providers (e.g., Keycloak) do not include + role claims in the UserInfo endpoint response. Args: access_token_str: The JWT access token string sso_jwt_handler: SSO-specific JWT handler for team ID extraction - result: The SSO result object to update with team IDs + result: The SSO result object to update with team IDs and role + role_mappings: Optional role mappings configuration for group-based role determination """ - if access_token_str and sso_jwt_handler and result: + if access_token_str and result: import jwt access_token_payload = jwt.decode( access_token_str, options={"verify_signature": False} ) - # Handle both dict and object result types - if isinstance(result, dict): - result_team_ids: Optional[List[str]] = result.get("team_ids", []) - if not result_team_ids: - team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) - result["team_ids"] = team_ids - else: - result_team_ids = getattr(result, "team_ids", []) if result else [] - if not result_team_ids: - team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) - setattr(result, "team_ids", team_ids) + # Extract team IDs from access token if sso_jwt_handler is available + if sso_jwt_handler: + if isinstance(result, dict): + result_team_ids: Optional[List[str]] = result.get("team_ids", []) + if not result_team_ids: + team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) + result["team_ids"] = team_ids + else: + result_team_ids = getattr(result, "team_ids", []) if result else [] + if not result_team_ids: + team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) + setattr(result, "team_ids", team_ids) + + # Extract user role from access token if not already set from UserInfo + existing_role = result.get("user_role") if isinstance(result, dict) else getattr(result, "user_role", None) + if existing_role is None: + user_role: Optional[LitellmUserRoles] = None + + # Try role_mappings first (group-based role determination) + if role_mappings is not None and role_mappings.roles: + group_claim = role_mappings.group_claim + user_groups_raw: Any = get_nested_value(access_token_payload, group_claim) + + user_groups: List[str] = [] + if isinstance(user_groups_raw, list): + user_groups = [str(g) for g in user_groups_raw] + elif isinstance(user_groups_raw, str): + user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] + elif user_groups_raw is not None: + user_groups = [str(user_groups_raw)] + + if user_groups: + user_role = determine_role_from_groups(user_groups, role_mappings) + verbose_proxy_logger.debug( + f"Determined role '{user_role}' from access token groups '{user_groups}' using role_mappings" + ) + elif role_mappings.default_role: + user_role = role_mappings.default_role + + # Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload + if user_role is None: + generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") + user_role_from_token = get_nested_value(access_token_payload, generic_user_role_attribute_name) + if user_role_from_token is not None: + user_role = get_litellm_user_role(user_role_from_token) + verbose_proxy_logger.debug( + f"Extracted role '{user_role}' from access token field '{generic_user_role_attribute_name}'" + ) + + if user_role is not None: + if isinstance(result, dict): + result["user_role"] = user_role + else: + setattr(result, "user_role", user_role) + verbose_proxy_logger.debug( + f"Set user_role='{user_role}' from JWT access token" + ) @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) @@ -688,7 +738,7 @@ async def get_generic_sso_response( ) access_token_str: Optional[str] = generic_sso.access_token - process_sso_jwt_access_token(access_token_str, sso_jwt_handler, result) + process_sso_jwt_access_token(access_token_str, sso_jwt_handler, result, role_mappings=role_mappings) except Exception as e: verbose_proxy_logger.exception( diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 41096503a2e..16f80826798 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -25,6 +25,8 @@ from litellm.proxy.management_endpoints.ui_sso import ( MicrosoftSSOHandler, SSOAuthenticationHandler, normalize_email, + process_sso_jwt_access_token, + determine_role_from_groups, _setup_team_mappings, ) from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -1298,6 +1300,7 @@ async def test_get_generic_sso_response_with_additional_headers(): # Mock the SSO provider and its methods mock_sso_instance = MagicMock() mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) + mock_sso_instance.access_token = None # Avoid triggering JWT decode in process_sso_jwt_access_token mock_sso_class = MagicMock(return_value=mock_sso_instance) @@ -1359,6 +1362,7 @@ async def test_get_generic_sso_response_with_empty_headers(): # Mock the SSO provider and its methods mock_sso_instance = MagicMock() mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) + mock_sso_instance.access_token = None # Avoid triggering JWT decode in process_sso_jwt_access_token mock_sso_class = MagicMock(return_value=mock_sso_instance) @@ -2546,22 +2550,25 @@ class TestProcessSSOJWTAccessToken: assert result.team_ids == [] def test_process_sso_jwt_access_token_no_sso_jwt_handler(self, sample_jwt_token): - """Test that nothing happens when sso_jwt_handler is None""" + """Test that JWT is decoded for role extraction even when sso_jwt_handler is None, + but team_ids are not extracted (team extraction requires sso_jwt_handler).""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, ) result = CustomOpenID(id="test_user", email="test@example.com", team_ids=[]) - with patch("jwt.decode") as mock_jwt_decode: + mock_payload = {"sub": "test_user", "email": "test@example.com"} + with patch("jwt.decode", return_value=mock_payload) as mock_jwt_decode: # Act process_sso_jwt_access_token( access_token_str=sample_jwt_token, sso_jwt_handler=None, result=result ) - # Assert nothing was processed - mock_jwt_decode.assert_not_called() + # JWT is decoded (for role extraction) but team_ids are not extracted + mock_jwt_decode.assert_called_once() assert result.team_ids == [] + assert result.user_role is None def test_process_sso_jwt_access_token_no_result( self, mock_jwt_handler, sample_jwt_token @@ -3848,3 +3855,219 @@ async def test_setup_team_mappings(): mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with( where={"id": "sso_config"} ) + + +# ============================================================================ +# Tests for get_litellm_user_role with list inputs (Keycloak returns lists) +# ============================================================================ + + +def test_get_litellm_user_role_with_string(): + """Test that get_litellm_user_role works with a plain string.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + result = get_litellm_user_role("proxy_admin") + assert result == LitellmUserRoles.PROXY_ADMIN + + +def test_get_litellm_user_role_with_list(): + """ + Test that get_litellm_user_role handles list inputs. + Keycloak returns roles as arrays like ["proxy_admin"] instead of strings. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + result = get_litellm_user_role(["proxy_admin"]) + assert result == LitellmUserRoles.PROXY_ADMIN + + +def test_get_litellm_user_role_with_empty_list(): + """Test that get_litellm_user_role returns None for empty lists.""" + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + result = get_litellm_user_role([]) + assert result is None + + +def test_get_litellm_user_role_with_invalid_role(): + """Test that get_litellm_user_role returns None for invalid roles.""" + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + result = get_litellm_user_role("not_a_real_role") + assert result is None + + +def test_get_litellm_user_role_with_list_multiple_roles(): + """Test that get_litellm_user_role takes the first element from a multi-element list.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + result = get_litellm_user_role(["proxy_admin", "internal_user"]) + assert result == LitellmUserRoles.PROXY_ADMIN + + +# ============================================================================ +# Tests for process_sso_jwt_access_token role extraction +# ============================================================================ + + +def test_process_sso_jwt_access_token_extracts_role_from_access_token(): + """ + Test that process_sso_jwt_access_token extracts user role from the JWT + access token when the UserInfo response did not include it. + + This is the core fix for the Keycloak SSO role mapping bug: Keycloak's + UserInfo endpoint does not return role claims, but the JWT access token + contains them. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + # Create a JWT access token with role claims (as Keycloak would) + access_token_payload = { + "sub": "user-123", + "email": "admin@test.com", + "litellm_role": ["proxy_admin"], + } + access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256") + + # Result object with no role set (simulating UserInfo response without roles) + result = CustomOpenID( + id="user-123", + email="admin@test.com", + display_name="Admin User", + team_ids=[], + user_role=None, + ) + + # Call with GENERIC_USER_ROLE_ATTRIBUTE pointing to litellm_role + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + +def test_process_sso_jwt_access_token_does_not_override_existing_role(): + """ + Test that process_sso_jwt_access_token does NOT override a role that was + already extracted from the UserInfo response. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + access_token_payload = { + "sub": "user-123", + "litellm_role": ["internal_user"], + } + access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256") + + # Result already has a role (e.g., set from UserInfo) + result = CustomOpenID( + id="user-123", + email="admin@test.com", + display_name="Admin User", + team_ids=[], + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + # Should keep the original role + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + +def test_process_sso_jwt_access_token_extracts_role_from_nested_field(): + """ + Test role extraction from a nested JWT field like resource_access.client.roles. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + access_token_payload = { + "sub": "user-123", + "resource_access": { + "my-client": { + "roles": ["proxy_admin"] + } + }, + } + access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256") + + result = CustomOpenID( + id="user-123", + email="admin@test.com", + display_name="Admin User", + team_ids=[], + user_role=None, + ) + + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "resource_access.my-client.roles"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + +def test_process_sso_jwt_access_token_with_role_mappings(): + """ + Test role extraction using role_mappings (group-based role determination) + from the JWT access token. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + + access_token_payload = { + "sub": "user-123", + "groups": ["keycloak-admins", "developers"], + } + access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256") + + result = CustomOpenID( + id="user-123", + email="admin@test.com", + display_name="Admin User", + team_ids=[], + user_role=None, + ) + + role_mappings = RoleMappings( + provider="generic", + group_claim="groups", + default_role=LitellmUserRoles.INTERNAL_USER, + roles={ + LitellmUserRoles.PROXY_ADMIN: ["keycloak-admins"], + LitellmUserRoles.INTERNAL_USER: ["developers"], + }, + ) + + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=role_mappings, + ) + + # Should get highest privilege role + assert result.user_role == LitellmUserRoles.PROXY_ADMIN From 0a55571f75cc6cad02aa3c5722ddafa5377053a8 Mon Sep 17 00:00:00 2001 From: Simon Sadedin Date: Sat, 7 Feb 2026 11:06:33 +1100 Subject: [PATCH 10/13] [Feat] add au version of `claude-opus-4-6` to model cost map (#20566) Notes: General support for Opus 4.6 was added in #20506 however it omitted the AU (australian) specific instance profile used in Bedrock. This change only adds the the au id. It is copied from the US model settings which is consistent with past additions of this regional model profile. --- model_prices_and_context_window.json | 30 ++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0da47634a94..4704549e716 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1263,6 +1263,36 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346 }, + "au.anthropic.claude-opus-4-6-v1:0": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_200k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1.1e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_200k_tokens": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "output_cost_per_token_above_200k_tokens": 4.125e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, From ba74e6d9d246b7b71d8ab71e2e3ef69e0ddfa748 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 6 Feb 2026 17:34:32 -0800 Subject: [PATCH 11/13] Add http support to custom code guardrails + Unified guardrails for MCP + Agent guardrail support (#20619) * fix: fix styling * fix(custom_code_guardrail.py): add http support for custom code guardrails allows users to call external guardrails on litellm with minimal code changes (no custom handlers) Test guardrail integrations more easily * feat(a2a/): add guardrails for agent interactions allows the same guardrails for llm's to be applied to agents as well * fix(a2a/): support passing guardrails to a2a from the UI * style(code-editor): allow editing custom code guardrails on ui + add examples of pre/post calls for custom code guardrails * feat(mcp/): support custom code guardrails for mcp calls allows custom code guardrails to work on mcp input * feat(chatui.tsx): support guardrails on mcp tool calls on playground --- .../proxy/guardrails/custom_code_guardrail.md | 58 ++- litellm/integrations/custom_guardrail.py | 1 + .../a2a/chat/guardrail_translation/README.md | 155 ++++++++ .../chat/guardrail_translation/__init__.py | 11 + .../a2a/chat/guardrail_translation/handler.py | 315 ++++++++++++++++ .../guardrail_translation/handler.py | 105 +++--- .../mcp_server/mcp_server_manager.py | 24 +- .../mcp_server/rest_endpoints.py | 87 +++-- .../proxy/agent_endpoints/a2a_endpoints.py | 33 +- litellm/proxy/common_request_processing.py | 16 +- .../custom_code/custom_code_guardrail.py | 33 +- .../guardrail_hooks/custom_code/primitives.py | 230 ++++++++++++ .../unified_guardrail/unified_guardrail.py | 2 + .../custom_code/CustomCodeModal.tsx | 277 +++++++++++--- .../components/guardrails/guardrail_info.tsx | 58 ++- .../mcp_tools/MCPToolArgumentsForm.tsx | 345 ++++++++++++++++++ .../src/components/networking.tsx | 20 +- .../components/playground/chat_ui/ChatUI.tsx | 324 ++++++++++++---- .../playground/chat_ui/chatConstants.ts | 1 + .../chat_ui/mode_endpoint_mapping.tsx | 1 + .../playground/llm_calls/a2a_send_message.tsx | 6 + 21 files changed, 1862 insertions(+), 240 deletions(-) create mode 100644 litellm/llms/a2a/chat/guardrail_translation/README.md create mode 100644 litellm/llms/a2a/chat/guardrail_translation/__init__.py create mode 100644 litellm/llms/a2a/chat/guardrail_translation/handler.py create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx diff --git a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md index cb246144497..8cbc247ae5e 100644 --- a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md +++ b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md @@ -61,15 +61,23 @@ curl -X POST http://localhost:4000/chat/completions \ ### Function Signature -Your code must define an `apply_guardrail` function: +Your code must define an `apply_guardrail` function. It can be either sync or async: ```python +# Sync version def apply_guardrail(inputs, request_data, input_type): # inputs: see table below # request_data: {"model": "...", "user_id": "...", "team_id": "...", "metadata": {...}} # input_type: "request" or "response" return allow() # or block() or modify() + +# Async version (recommended when using HTTP primitives) +async def apply_guardrail(inputs, request_data, input_type): + response = await http_post("https://api.example.com/check", body={"text": inputs["texts"][0]}) + if response["success"] and response["body"].get("flagged"): + return block("Content flagged") + return allow() ``` ### `inputs` Parameter @@ -145,6 +153,29 @@ def apply_guardrail(inputs, request_data, input_type): | `char_count(text)` | Count characters | | `lower(text)` / `upper(text)` / `trim(text)` | String transforms | +### HTTP Requests (Async) + +Make async HTTP requests to external APIs for additional validation or content moderation. + +| Function | Description | +|----------|-------------| +| `await http_request(url, method, headers, body, timeout)` | General async HTTP request | +| `await http_get(url, headers, timeout)` | Async GET request | +| `await http_post(url, body, headers, timeout)` | Async POST request | + +**Response format:** +```python +{ + "status_code": 200, # HTTP status code + "body": {...}, # Response body (parsed JSON or string) + "headers": {...}, # Response headers + "success": True, # True if status code is 2xx + "error": None # Error message if request failed +} +``` + +**Note:** When using HTTP primitives, define your function as `async def apply_guardrail(...)` for non-blocking execution. + ## Examples ### Block PII (SSN) @@ -213,6 +244,29 @@ def apply_guardrail(inputs, request_data, input_type): return allow() ``` +### Call External Moderation API (Async) + +```python +async def apply_guardrail(inputs, request_data, input_type): + # Call an external moderation API + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text, "user_id": request_data["user_id"]}, + headers={"Authorization": "Bearer YOUR_API_KEY"}, + timeout=10 + ) + + if not response["success"]: + # API call failed - decide whether to allow or block + return allow() + + if response["body"].get("flagged"): + return block(response["body"].get("reason", "Content flagged")) + + return allow() +``` + ### Combine Multiple Checks ```python @@ -241,8 +295,8 @@ Custom code runs in a restricted environment: - ❌ No `import` statements - ❌ No file I/O -- ❌ No network access - ❌ No `exec()` or `eval()` +- ✅ HTTP requests via built-in `http_request`, `http_get`, `http_post` primitives - ✅ Only LiteLLM-provided primitives available ## Per-Request Usage diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1652ec2aa0c..bbd55a59bce 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -268,6 +268,7 @@ class CustomGuardrail(CustomLogger): """ Returns the guardrail(s) to be run from the metadata or root """ + if "guardrails" in data: return data["guardrails"] metadata = data.get("litellm_metadata") or data.get("metadata", {}) diff --git a/litellm/llms/a2a/chat/guardrail_translation/README.md b/litellm/llms/a2a/chat/guardrail_translation/README.md new file mode 100644 index 00000000000..1e18f5cda3a --- /dev/null +++ b/litellm/llms/a2a/chat/guardrail_translation/README.md @@ -0,0 +1,155 @@ +# A2A Protocol Guardrail Translation Handler + +Handler for processing A2A (Agent-to-Agent) Protocol messages with guardrails. + +## Overview + +This handler processes A2A JSON-RPC 2.0 input/output by: +1. Extracting text from message parts (`kind: "text"`) +2. Applying guardrails to text content +3. Mapping guardrailed text back to original structure + +## A2A Protocol Format + +### Input Format (JSON-RPC 2.0) + +```json +{ + "jsonrpc": "2.0", + "id": "request-id", + "method": "message/send", + "params": { + "message": { + "kind": "message", + "messageId": "...", + "role": "user", + "parts": [ + {"kind": "text", "text": "Hello, my SSN is 123-45-6789"} + ] + }, + "metadata": { + "guardrails": ["block-ssn"] + } + } +} +``` + +### Output Formats + +The handler supports multiple A2A response formats: + +**Direct message:** +```json +{ + "result": { + "kind": "message", + "parts": [{"kind": "text", "text": "Response text"}] + } +} +``` + +**Nested message:** +```json +{ + "result": { + "message": { + "parts": [{"kind": "text", "text": "Response text"}] + } + } +} +``` + +**Task with artifacts:** +```json +{ + "result": { + "kind": "task", + "artifacts": [ + {"parts": [{"kind": "text", "text": "Artifact text"}]} + ] + } +} +``` + +**Task with status message:** +```json +{ + "result": { + "kind": "task", + "status": { + "message": { + "parts": [{"kind": "text", "text": "Status message"}] + } + } + } +} +``` + +**Streaming artifact-update:** +```json +{ + "result": { + "kind": "artifact-update", + "artifact": { + "parts": [{"kind": "text", "text": "Streaming text"}] + } + } +} +``` + +## Usage + +The handler is automatically discovered and applied when guardrails are used with A2A endpoints. + +### Via LiteLLM Proxy + +```bash +curl -X POST 'http://localhost:4000/a2a/my-agent' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "jsonrpc": "2.0", + "id": "1", + "method": "message/send", + "params": { + "message": { + "kind": "message", + "messageId": "msg-1", + "role": "user", + "parts": [{"kind": "text", "text": "Hello, my SSN is 123-45-6789"}] + }, + "metadata": { + "guardrails": ["block-ssn"] + } + } +}' +``` + +### Specifying Guardrails + +Guardrails can be specified in the A2A request via the `metadata.guardrails` field: + +```json +{ + "params": { + "message": {...}, + "metadata": { + "guardrails": ["block-ssn", "pii-filter"] + } + } +} +``` + +## Extension + +Override these methods to customize behavior: + +- `_extract_texts_from_result()`: Custom text extraction from A2A responses +- `_extract_texts_from_parts()`: Custom text extraction from message parts +- `_apply_text_to_path()`: Custom application of guardrailed text + +## Call Types + +This handler is registered for: +- `CallTypes.send_message`: Synchronous A2A message sending +- `CallTypes.asend_message`: Asynchronous A2A message sending diff --git a/litellm/llms/a2a/chat/guardrail_translation/__init__.py b/litellm/llms/a2a/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..13c20677485 --- /dev/null +++ b/litellm/llms/a2a/chat/guardrail_translation/__init__.py @@ -0,0 +1,11 @@ +"""A2A Protocol handler for Unified Guardrails.""" + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.send_message: A2AGuardrailHandler, + CallTypes.asend_message: A2AGuardrailHandler, +} + +__all__ = ["guardrail_translation_mappings"] diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py new file mode 100644 index 00000000000..770453f2def --- /dev/null +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -0,0 +1,315 @@ +""" +A2A Protocol Handler for Unified Guardrails + +This module provides guardrail translation support for A2A (Agent-to-Agent) Protocol. +It handles both JSON-RPC 2.0 input requests and output responses, extracting text +from message parts and applying guardrails. + +A2A Protocol Format: +- Input: JSON-RPC 2.0 with params.message.parts containing text parts +- Output: JSON-RPC 2.0 with result containing message/artifact parts +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth + + +class A2AGuardrailHandler(BaseTranslation): + """ + Handler for processing A2A Protocol messages with guardrails. + + This class provides methods to: + 1. Process input messages (pre-call hook) - extracts text from A2A message parts + 2. Process output responses (post-call hook) - extracts text from A2A response parts + + A2A Message Format: + - Input: params.message.parts[].text (where kind == "text") + - Output: result.message.parts[].text or result.artifacts[].parts[].text + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> Any: + """ + Process A2A input messages by applying guardrails to text content. + + Extracts text from A2A message parts and applies guardrails. + + Args: + data: The A2A JSON-RPC 2.0 request data + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + + Returns: + Modified data with guardrails applied to text content + """ + # A2A request format: { "params": { "message": { "parts": [...] } } } + params = data.get("params", {}) + message = params.get("message", {}) + parts = message.get("parts", []) + + if not parts: + verbose_proxy_logger.debug("A2A: No parts in message, skipping guardrail") + return data + + texts_to_check: List[str] = [] + text_part_indices: List[int] = [] # Track which parts contain text + + # Step 1: Extract text from all text parts + for part_idx, part in enumerate(parts): + if part.get("kind") == "text": + text = part.get("text", "") + if text: + texts_to_check.append(text) + text_part_indices.append(part_idx) + + # Step 2: Apply guardrail to all texts in batch + if texts_to_check: + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + + # Pass the structured A2A message to guardrails + inputs["structured_messages"] = [message] + + # Include agent model info if available + model = data.get("model") + if model: + inputs["model"] = model + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + + guardrailed_texts = guardrailed_inputs.get("texts", []) + + # Step 3: Apply guardrailed text back to original parts + if guardrailed_texts and len(guardrailed_texts) == len(text_part_indices): + for task_idx, part_idx in enumerate(text_part_indices): + parts[part_idx]["text"] = guardrailed_texts[task_idx] + + verbose_proxy_logger.debug("A2A: Processed input message: %s", message) + + return data + + async def process_output_response( + self, + response: Any, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + ) -> Any: + """ + Process A2A output response by applying guardrails to text content. + + Handles multiple A2A response formats: + - Direct message: {"result": {"kind": "message", "parts": [...]}} + - Nested message: {"result": {"message": {"parts": [...]}}} + - Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}} + - Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}} + + Args: + response: A2A JSON-RPC 2.0 response dict or object + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + user_api_key_dict: User API key metadata + + Returns: + Modified response with guardrails applied to text content + """ + # Handle both dict and Pydantic model responses + if hasattr(response, "model_dump"): + response_dict = response.model_dump() + is_pydantic = True + elif isinstance(response, dict): + response_dict = response + is_pydantic = False + else: + verbose_proxy_logger.warning( + "A2A: Unknown response type %s, skipping guardrail", type(response) + ) + return response + + result = response_dict.get("result", {}) + if not result or not isinstance(result, dict): + verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail") + return response + + # Find all text-containing parts in the response + texts_to_check: List[str] = [] + # Each mapping is (path_to_parts_list, part_index) + # path_to_parts_list is a tuple of keys to navigate to the parts list + task_mappings: List[Tuple[Tuple[str, ...], int]] = [] + + # Extract texts from all possible locations + self._extract_texts_from_result( + result=result, + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + if not texts_to_check: + verbose_proxy_logger.debug("A2A: No text content in response") + return response + + # Step 2: Apply guardrail to all texts in batch + # Create a request_data dict with response info and user API key metadata + request_data: dict = {"response": response_dict} + + # Add user API key metadata with prefixed keys + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + + guardrailed_texts = guardrailed_inputs.get("texts", []) + + # Step 3: Apply guardrailed text back to original response + if guardrailed_texts and len(guardrailed_texts) == len(task_mappings): + for task_idx, (path, part_idx) in enumerate(task_mappings): + self._apply_text_to_path( + result=result, + path=path, + part_idx=part_idx, + text=guardrailed_texts[task_idx], + ) + + verbose_proxy_logger.debug("A2A: Processed output response") + + # Update the original response + if is_pydantic: + # For Pydantic models, we need to update the underlying dict + # and the model will reflect the changes + response_dict["result"] = result + return response + else: + response["result"] = result + return response + + def _extract_texts_from_result( + self, + result: Dict[str, Any], + texts_to_check: List[str], + task_mappings: List[Tuple[Tuple[str, ...], int]], + ) -> None: + """ + Extract text from all possible locations in an A2A result. + + Handles multiple response formats: + 1. Direct message with parts: {"parts": [...]} + 2. Nested message: {"message": {"parts": [...]}} + 3. Task with artifacts: {"artifacts": [{"parts": [...]}]} + 4. Task with status message: {"status": {"message": {"parts": [...]}}} + 5. Streaming artifact-update: {"artifact": {"parts": [...]}} + """ + # Case 1: Direct parts in result (direct message) + if "parts" in result: + self._extract_texts_from_parts( + parts=result["parts"], + path=("parts",), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 2: Nested message + message = result.get("message") + if message and isinstance(message, dict) and "parts" in message: + self._extract_texts_from_parts( + parts=message["parts"], + path=("message", "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 3: Streaming artifact-update (singular artifact) + artifact = result.get("artifact") + if artifact and isinstance(artifact, dict) and "parts" in artifact: + self._extract_texts_from_parts( + parts=artifact["parts"], + path=("artifact", "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 4: Task with status message + status = result.get("status", {}) + if isinstance(status, dict): + status_message = status.get("message") + if ( + status_message + and isinstance(status_message, dict) + and "parts" in status_message + ): + self._extract_texts_from_parts( + parts=status_message["parts"], + path=("status", "message", "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 5: Task with artifacts (plural, array) + artifacts = result.get("artifacts", []) + if artifacts and isinstance(artifacts, list): + for artifact_idx, art in enumerate(artifacts): + if isinstance(art, dict) and "parts" in art: + self._extract_texts_from_parts( + parts=art["parts"], + path=("artifacts", str(artifact_idx), "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + def _extract_texts_from_parts( + self, + parts: List[Dict[str, Any]], + path: Tuple[str, ...], + texts_to_check: List[str], + task_mappings: List[Tuple[Tuple[str, ...], int]], + ) -> None: + """Extract text from message parts.""" + for part_idx, part in enumerate(parts): + if part.get("kind") == "text": + text = part.get("text", "") + if text: + texts_to_check.append(text) + task_mappings.append((path, part_idx)) + + def _apply_text_to_path( + self, + result: Dict[Union[str, int], Any], + path: Tuple[str, ...], + part_idx: int, + text: str, + ) -> None: + """Apply guardrailed text back to the specified path in the result.""" + # Navigate to the parts list + current = result + for key in path: + if key.isdigit(): + # Array index + current = current[int(key)] + else: + current = current[key] + + # Update the text in the part + current[part_idx]["text"] = text diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 4d53ae7059d..14bbb82808d 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -1,26 +1,37 @@ """ MCP Guardrail Handler for Unified Guardrails. -This handler works with the synthetic "messages" payload generated by -`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user -message whose `content` string encodes the MCP tool name and arguments. The -handler simply feeds that text through the configured guardrail and writes the -result back onto the message. +Converts an MCP call_tool (name + arguments) into a single OpenAI-compatible +tool_call and passes it to apply_guardrail. Works with the synthetic payload +from ProxyLogging._convert_mcp_to_llm_format. + +Note: For MCP tool definitions (schema) -> OpenAI tools=[], see +litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool +when you have a full MCP Tool from list_tools. Here we only have the call +payload (name + arguments) so we just build the tool_call. """ from typing import TYPE_CHECKING, Any, Dict, Optional +from mcp.types import Tool as MCPTool + from litellm._logging import verbose_proxy_logger +from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.llms.openai import ( + ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail from mcp.types import CallToolResult + from litellm.integrations.custom_guardrail import CustomGuardrail + class MCPGuardrailTranslationHandler(BaseTranslation): - """Guardrail translation handler for MCP tool calls.""" + """Guardrail translation handler for MCP tool calls (passes a single tool_call to guardrail).""" async def process_input_messages( self, @@ -28,56 +39,51 @@ class MCPGuardrailTranslationHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, ) -> Dict[str, Any]: - messages = data.get("messages") - if not isinstance(messages, list) or not messages: - verbose_proxy_logger.debug("MCP Guardrail: No messages to process") + mcp_tool_name = data.get("mcp_tool_name") or data.get("name") + mcp_arguments = data.get("mcp_arguments") or data.get("arguments") + mcp_tool_description = data.get("mcp_tool_description") or data.get( + "description" + ) + if mcp_arguments is None or not isinstance(mcp_arguments, dict): + mcp_arguments = {} + + if not mcp_tool_name: + verbose_proxy_logger.debug("MCP Guardrail: mcp_tool_name missing") return data - first_message = messages[0] - content: Optional[str] = None - if isinstance(first_message, dict): - content = first_message.get("content") - else: - content = getattr(first_message, "content", None) + # Convert MCP input via transform_mcp_tool_to_openai_tool, then map to litellm + # ChatCompletionToolParam (openai SDK type has incompatible strict/cache_control). + mcp_tool = MCPTool( + name=mcp_tool_name, + description=mcp_tool_description or "", + inputSchema={}, # Call payload has no schema; guardrail gets args from request_data + ) + openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool) + fn = openai_tool["function"] + tool_def: ChatCompletionToolParam = { + "type": "function", + "function": ChatCompletionToolParamFunctionChunk( + name=fn["name"], + description=fn.get("description") or "", + parameters=fn.get("parameters") + or { + "type": "object", + "properties": {}, + "additionalProperties": False, + }, + strict=fn.get("strict", False) or False, # Default to False if None + ), + } + inputs: GenericGuardrailAPIInputs = GenericGuardrailAPIInputs( + tools=[tool_def], + ) - if not isinstance(content, str): - verbose_proxy_logger.debug( - "MCP Guardrail: Message content missing or not a string", - ) - return data - - inputs = GenericGuardrailAPIInputs(texts=[content]) - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) - guardrailed_texts = ( - guardrailed_inputs.get("texts", []) if guardrailed_inputs else [] - ) - - if guardrailed_texts: - new_content = guardrailed_texts[0] - if isinstance(first_message, dict): - first_message["content"] = new_content - else: - setattr(first_message, "content", new_content) - - verbose_proxy_logger.debug( - "MCP Guardrail: Updated content for tool %s", - data.get("mcp_tool_name"), - ) - else: - verbose_proxy_logger.debug( - "MCP Guardrail: Guardrail returned no text updates for tool %s", - data.get("mcp_tool_name"), - ) - return data async def process_output_response( @@ -87,7 +93,6 @@ class MCPGuardrailTranslationHandler(BaseTranslation): litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, ) -> Any: - # Not implemented: MCP guardrail translation never calls this path today. verbose_proxy_logger.debug( "MCP Guardrail: Output processing not implemented for MCP tools", ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4c17a2ff3e0..f5b0f152fd8 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -65,12 +65,13 @@ from litellm.types.mcp_server.mcp_server_manager import ( from litellm.types.utils import CallTypes try: - from mcp.shared.tool_name_validation import ( # type: ignore - SEP_986_URL, - validate_tool_name, + from mcp.shared.tool_name_validation import ( + validate_tool_name, # type: ignore[reportAssignmentType] ) + from mcp.shared.tool_name_validation import SEP_986_URL except ImportError: from pydantic import BaseModel + SEP_986_URL = "https://github.com/modelcontextprotocol/protocol/blob/main/proposals/0001-tool-name-validation.md" class ToolNameValidationResult(BaseModel): @@ -469,12 +470,12 @@ class MCPServerManager: ) # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[ - base_tool_name - ] = server_prefix - self.tool_name_to_mcp_server_name_mapping[ - prefixed_tool_name - ] = server_prefix + self.tool_name_to_mcp_server_name_mapping[base_tool_name] = ( + server_prefix + ) + self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = ( + server_prefix + ) registered_count += 1 verbose_logger.debug( @@ -1929,7 +1930,9 @@ class MCPServerManager: ) async def _call_tool_via_client(client, params): - return await client.call_tool(params, host_progress_callback=host_progress_callback) + return await client.call_tool( + params, host_progress_callback=host_progress_callback + ) tasks.append( asyncio.create_task(_call_tool_via_client(client, call_tool_params)) @@ -1967,7 +1970,6 @@ class MCPServerManager: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, host_progress_callback: Optional[Callable] = None, - ) -> CallToolResult: """ Call a tool with the given name and arguments diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d93f852f22d..eb47fb4f608 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -12,6 +12,7 @@ from litellm.proxy._experimental.mcp_server.utils import merge_mcp_headers from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.mcp import MCPAuth +from litellm.types.utils import CallTypes MCP_AVAILABLE: bool = True try: @@ -28,6 +29,7 @@ router = APIRouter( if MCP_AVAILABLE: from mcp.types import Tool as MCPTool + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -96,6 +98,35 @@ if MCP_AVAILABLE: return _create_tool_response_objects(tools, server.mcp_info) + async def _resolve_allowed_mcp_servers_for_tool_call( + user_api_key_dict: UserAPIKeyAuth, + server_id: str, + ) -> List[MCPServer]: + """Resolve allowed MCP servers for the given user and validate server_id access.""" + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + allowed_server_ids_set = set() + for auth_context in auth_contexts: + servers = await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth=auth_context + ) + allowed_server_ids_set.update(servers) + if server_id not in allowed_server_ids_set: + raise HTTPException( + status_code=403, + detail={ + "error": "access_denied", + "message": f"The key is not allowed to access server {server_id}", + }, + ) + allowed_mcp_servers: List[MCPServer] = [] + for allowed_server_id in allowed_server_ids_set: + server = global_mcp_server_manager.get_mcp_server_by_id( + allowed_server_id + ) + if server is not None: + allowed_mcp_servers.append(server) + return allowed_mcp_servers + ######################################################## @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) async def list_tool_rest_api( @@ -261,7 +292,14 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) - from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + from litellm.proxy.proxy_server import ( + general_settings, + proxy_config, + proxy_logging_obj, + ) try: data = await request.json() @@ -289,11 +327,16 @@ if MCP_AVAILABLE: tool_arguments = data.get("arguments") - data = await add_litellm_data_to_request( - data=data, - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, + proxy_base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + data, logging_obj = ( + await proxy_base_llm_response_processor.common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ) ) # FIX: Extract MCP auth headers from request @@ -322,35 +365,9 @@ if MCP_AVAILABLE: if "metadata" in data and "user_api_key_auth" in data["metadata"]: data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] - # Get all auth contexts - auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - - # Collect allowed server IDs from all contexts - allowed_server_ids_set = set() - for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth=auth_context - ) - allowed_server_ids_set.update(servers) - - # Check if the specified server_id is allowed - if server_id not in allowed_server_ids_set: - raise HTTPException( - status_code=403, - detail={ - "error": "access_denied", - "message": f"The key is not allowed to access server {server_id}", - }, - ) - - # Build allowed_mcp_servers list (only include allowed servers) - allowed_mcp_servers: List[MCPServer] = [] - for allowed_server_id in allowed_server_ids_set: - server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_server_id - ) - if server is not None: - allowed_mcp_servers.append(server) + allowed_mcp_servers = await _resolve_allowed_mcp_servers_for_tool_call( + user_api_key_dict, server_id + ) # Call execute_mcp_tool directly (permission checks already done) result = await execute_mcp_tool( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 24727aacd75..12b2d5c4dfe 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -14,6 +14,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.utils import all_litellm_params router = APIRouter() @@ -75,10 +76,7 @@ async def _handle_stream_message( return StreamingResponse(_error_stream(), media_type="application/x-ndjson") - from a2a.types import ( - MessageSendParams, - SendStreamingMessageRequest, - ) + from a2a.types import MessageSendParams, SendStreamingMessageRequest async def stream_response(): try: @@ -208,16 +206,17 @@ async def invoke_agent_a2a( from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ( general_settings, proxy_config, + proxy_logging_obj, version, ) body = {} try: body = await request.json() + verbose_proxy_logger.debug(f"A2A request for agent '{agent_id}': {body}") # Validate JSON-RPC format @@ -230,6 +229,16 @@ async def invoke_agent_a2a( method = body.get("method") params = body.get("params", {}) + if params: + # extract any litellm params from the params - eg. 'guardrails' + params_to_remove = [] + for key, value in params.items(): + if key in all_litellm_params: + params_to_remove.append(key) + body[key] = value + for key in params_to_remove: + params.pop(key) + if not A2A_SDK_AVAILABLE: return _jsonrpc_error( request_id, @@ -283,12 +292,18 @@ async def invoke_agent_a2a( ) # Add litellm data (user_api_key, user_id, team_id, etc.) - data = await add_litellm_data_to_request( - data=body, + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + processor = ProxyBaseLLMRequestProcessing(data=body) + data, logging_obj = await processor.common_processing_pre_call_logic( request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type="asend_message", version=version, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3e079045774..f33b2412260 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -267,7 +267,9 @@ def _override_openai_response_model( hidden_params = getattr(response_obj, "_hidden_params", {}) or {} if isinstance(hidden_params, dict): fallback_headers = hidden_params.get("additional_headers", {}) or {} - attempted_fallbacks = fallback_headers.get("x-litellm-attempted-fallbacks", None) + attempted_fallbacks = fallback_headers.get( + "x-litellm-attempted-fallbacks", None + ) if attempted_fallbacks is not None and attempted_fallbacks > 0: # A fallback occurred - preserve the actual model that was used verbose_proxy_logger.debug( @@ -517,6 +519,8 @@ class ProxyBaseLLMRequestProcessing: "aget_interaction", "adelete_interaction", "acancel_interaction", + "asend_message", + "call_mcp_tool", ], version: Optional[str] = None, user_model: Optional[str] = None, @@ -837,7 +841,9 @@ class ProxyBaseLLMRequestProcessing: # aliasing/routing, but the OpenAI-compatible response `model` field should reflect # what the client sent. if requested_model_from_client: - self.data["_litellm_client_requested_model"] = requested_model_from_client + self.data["_litellm_client_requested_model"] = ( + requested_model_from_client + ) if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): @@ -1409,9 +1415,9 @@ class ProxyBaseLLMRequestProcessing: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs[ - "cache_creation_input_tokens" - ] = cache_creation_input_tokens + usage_kwargs["cache_creation_input_tokens"] = ( + cache_creation_input_tokens + ) if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index a0ca324411c..68f9dfd7abc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -5,7 +5,7 @@ This module provides a guardrail that executes user-defined Python-like code to implement custom guardrail logic. The code runs in a sandboxed environment with access to LiteLLM-provided primitives for common guardrail operations. -Example custom code: +Example custom code (sync): def apply_guardrail(inputs, request_data, input_type): '''Block messages containing SSNs''' @@ -13,8 +13,22 @@ Example custom code: if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): return block("Social Security Number detected") return allow() + +Example custom code (async with HTTP): + + async def apply_guardrail(inputs, request_data, input_type): + '''Call external moderation API''' + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text} + ) + if response["success"] and response["body"].get("flagged"): + return block("Content flagged by moderation API") + return allow() """ +import asyncio import threading from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast @@ -101,6 +115,9 @@ class CustomCodeGuardrail(CustomGuardrail): GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.during_mcp_call, + GuardrailEventHooks.logging_only, ] super().__init__( @@ -175,6 +192,13 @@ class CustomCodeGuardrail(CustomGuardrail): This method calls the user-defined apply_guardrail function and processes its result to determine the appropriate action. + The user-defined function can be either sync or async: + - Sync: def apply_guardrail(inputs, request_data, input_type): ... + - Async: async def apply_guardrail(inputs, request_data, input_type): ... + + Async functions are recommended when using http_request, http_get, or + http_post primitives to avoid blocking the event loop. + Args: inputs: Dictionary containing texts, images, tool_calls request_data: The original request data with metadata @@ -188,6 +212,7 @@ class CustomCodeGuardrail(CustomGuardrail): HTTPException: If content is blocked CustomCodeExecutionError: If execution fails """ + if self._compiled_function is None: if self._compile_error: raise CustomCodeExecutionError( @@ -201,9 +226,13 @@ class CustomCodeGuardrail(CustomGuardrail): # Prepare request_data with safe subset of information safe_request_data = self._prepare_safe_request_data(request_data) - # Execute the custom function + # Execute the custom function - handle both sync and async functions result = self._compiled_function(inputs, safe_request_data, input_type) + # If the function is async (returns a coroutine), await it + if asyncio.iscoroutine(result): + result = await result + # Process the result return self._process_result( result=result, diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 695e59977c8..de7690635d8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -10,7 +10,11 @@ import re from typing import Any, Dict, List, Optional, Tuple, Type, Union from urllib.parse import urlparse +import httpx + from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider # ============================================================================= # Result Types - Used by Starlark code to return guardrail decisions @@ -349,6 +353,228 @@ def get_url_domain(url: str) -> Optional[str]: return None +# ============================================================================= +# HTTP Request Primitives (Async) +# ============================================================================= + +# Default timeout for HTTP requests (in seconds) +_HTTP_DEFAULT_TIMEOUT = 30.0 + +# Maximum allowed timeout (in seconds) +_HTTP_MAX_TIMEOUT = 60.0 + + +def _http_error_response(error: str) -> Dict[str, Any]: + """Create a standardized error response for HTTP requests.""" + return { + "status_code": 0, + "body": None, + "headers": {}, + "success": False, + "error": error, + } + + +def _http_success_response(response: httpx.Response) -> Dict[str, Any]: + """Create a standardized success response from an httpx Response.""" + parsed_body: Any + try: + parsed_body = response.json() + except (json.JSONDecodeError, ValueError): + parsed_body = response.text + + return { + "status_code": response.status_code, + "body": parsed_body, + "headers": dict(response.headers), + "success": 200 <= response.status_code < 300, + "error": None, + } + + +def _prepare_http_body( + body: Optional[Any], +) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Prepare body arguments for HTTP request - returns (json_body, data_body).""" + if body is None: + return None, None + if isinstance(body, dict): + return body, None + if isinstance(body, list): + return None, json.dumps(body) + if isinstance(body, str): + return None, body + return None, str(body) + + +async def http_request( + url: str, + method: str = "GET", + headers: Optional[Dict[str, str]] = None, + body: Optional[Any] = None, + timeout: Optional[float] = None, +) -> Dict[str, Any]: + """ + Make an async HTTP request to an external service. + + This function allows custom guardrails to call external APIs for + additional validation, content moderation, or data enrichment. + + Uses LiteLLM's global cached AsyncHTTPHandler for connection pooling + and better performance. + + Args: + url: The URL to request + method: HTTP method (GET, POST, PUT, DELETE, PATCH). Defaults to GET. + headers: Optional dict of HTTP headers + body: Optional request body (will be JSON-encoded if dict/list) + timeout: Optional timeout in seconds (default: 30, max: 60) + + Returns: + Dict containing: + - status_code: HTTP status code + - body: Response body (parsed as JSON if possible, otherwise string) + - headers: Response headers as dict + - success: True if status code is 2xx + - error: Error message if request failed, None otherwise + + Example: + # Simple GET request + response = await http_request("https://api.example.com/check") + if response["success"]: + data = response["body"] + + # POST request with JSON body + response = await http_request( + "https://api.example.com/moderate", + method="POST", + headers={"Authorization": "Bearer token"}, + body={"text": "content to check"} + ) + """ + # Validate URL + if not is_valid_url(url): + return _http_error_response(f"Invalid URL: {url}") + + # Validate and normalize method + method = method.upper() + allowed_methods = {"GET", "POST", "PUT", "DELETE", "PATCH"} + if method not in allowed_methods: + return _http_error_response( + f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}" + ) + + # Apply timeout limits + if timeout is None: + timeout = _HTTP_DEFAULT_TIMEOUT + else: + timeout = min(max(0.1, timeout), _HTTP_MAX_TIMEOUT) + + # Get the global cached async HTTP client + client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)}, + ) + + try: + response = await _execute_http_request( + client, method, url, headers, body, timeout + ) + return _http_success_response(response) + + except httpx.TimeoutException as e: + verbose_proxy_logger.warning(f"Custom code http_request timeout: {e}") + return _http_error_response(f"Request timeout after {timeout}s") + except httpx.HTTPStatusError as e: + # Return the response even for non-2xx status codes + return _http_success_response(e.response) + except httpx.RequestError as e: + verbose_proxy_logger.warning(f"Custom code http_request error: {e}") + return _http_error_response(f"Request failed: {str(e)}") + except Exception as e: + verbose_proxy_logger.warning(f"Custom code http_request unexpected error: {e}") + return _http_error_response(f"Unexpected error: {str(e)}") + + +async def _execute_http_request( + client: Any, + method: str, + url: str, + headers: Optional[Dict[str, str]], + body: Optional[Any], + timeout: float, +) -> httpx.Response: + """Execute the HTTP request using the appropriate client method.""" + json_body, data_body = _prepare_http_body(body) + + if method == "GET": + return await client.get(url=url, headers=headers) + elif method == "POST": + return await client.post( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + elif method == "PUT": + return await client.put( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + elif method == "DELETE": + return await client.delete( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + elif method == "PATCH": + return await client.patch( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + +async def http_get( + url: str, + headers: Optional[Dict[str, str]] = None, + timeout: Optional[float] = None, +) -> Dict[str, Any]: + """ + Make an async HTTP GET request. + + Convenience wrapper around http_request for GET requests. + + Args: + url: The URL to request + headers: Optional dict of HTTP headers + timeout: Optional timeout in seconds + + Returns: + Same as http_request + """ + return await http_request(url=url, method="GET", headers=headers, timeout=timeout) + + +async def http_post( + url: str, + body: Optional[Any] = None, + headers: Optional[Dict[str, str]] = None, + timeout: Optional[float] = None, +) -> Dict[str, Any]: + """ + Make an async HTTP POST request. + + Convenience wrapper around http_request for POST requests. + + Args: + url: The URL to request + body: Optional request body (will be JSON-encoded if dict/list) + headers: Optional dict of HTTP headers + timeout: Optional timeout in seconds + + Returns: + Same as http_request + """ + return await http_request( + url=url, method="POST", headers=headers, body=body, timeout=timeout + ) + + # ============================================================================= # Code Detection Primitives # ============================================================================= @@ -575,6 +801,10 @@ def get_custom_code_primitives() -> Dict[str, Any]: "is_valid_url": is_valid_url, "all_urls_valid": all_urls_valid, "get_url_domain": get_url_domain, + # HTTP (async) + "http_request": http_request, + "http_get": http_get, + "http_post": http_post, # Code detection "detect_code": detect_code, "detect_code_languages": detect_code_languages, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index f07f65d10f5..cc05358baf7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -51,6 +51,7 @@ class UnifiedLLMGuardrails(CustomLogger): Runs on only Input Use this if you want to MODIFY the input """ + global endpoint_guardrail_translation_mappings from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -66,6 +67,7 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.pre_mcp_call + if ( guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True diff --git a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx index 73d96c590ec..1b8a8c034bd 100644 --- a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx @@ -9,7 +9,7 @@ import { CaretRightOutlined, SaveOutlined, } from "@ant-design/icons"; -import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking"; +import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; const { Panel } = Collapse; @@ -19,7 +19,7 @@ const { TextArea } = Input; const CODE_TEMPLATES = { empty: { name: "Empty Template", - code: `def apply_guardrail(inputs, request_data, input_type): + code: `async def apply_guardrail(inputs, request_data, input_type): # inputs: {texts, images, tools, tool_calls, structured_messages, model} # request_data: {model, user_id, team_id, end_user_id, metadata} # input_type: "request" or "response" @@ -68,6 +68,27 @@ const CODE_TEMPLATES = { return block("Response missing required fields") return allow()`, }, + externalAPI: { + name: "External API Check (async)", + code: `async def apply_guardrail(inputs, request_data, input_type): + # Call an external moderation API (async for non-blocking) + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text, "user_id": request_data["user_id"]}, + headers={"Authorization": "Bearer YOUR_API_KEY"}, + timeout=10 + ) + + if not response["success"]: + # API call failed, allow by default or block + return allow() + + if response["body"].get("flagged"): + return block(response["body"].get("reason", "Content flagged")) + + return allow()`, + }, }; // Available primitives organized by category @@ -77,6 +98,11 @@ const PRIMITIVES = { { name: "block(reason)", desc: "Reject with message" }, { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, ], + "HTTP Requests (async)": [ + { name: "await http_request(url, method, headers, body)", desc: "Make async HTTP request" }, + { name: "await http_get(url, headers)", desc: "Async GET request" }, + { name: "await http_post(url, body, headers)", desc: "Async POST request" }, + ], "Regex Functions": [ { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" }, { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" }, @@ -111,13 +137,30 @@ const MODE_OPTIONS = [ { value: "post_call", label: "post_call (Response)" }, { value: "during_call", label: "during_call (Parallel)" }, { value: "logging_only", label: "logging_only" }, + { value: "pre_mcp_call", label: "pre_mcp_call (Before MCP Tool Call)" }, + { value: "post_mcp_call", label: "post_mcp_call (After MCP Tool Call)" }, + { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" }, ]; +// Data for editing an existing guardrail +export interface EditGuardrailData { + guardrail_id: string; + guardrail_name: string; + litellm_params: { + mode?: string | string[]; + default_on?: boolean; + custom_code?: string; + [key: string]: any; + }; +} + interface CustomCodeModalProps { visible: boolean; onClose: () => void; onSuccess: () => void; accessToken: string | null; + /** If provided, the modal will be in edit mode */ + editData?: EditGuardrailData | null; } const CustomCodeModal: React.FC = ({ @@ -125,16 +168,72 @@ const CustomCodeModal: React.FC = ({ onClose, onSuccess, accessToken, + editData, }) => { + const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); - const [mode, setMode] = useState("pre_call"); + const [mode, setMode] = useState(["pre_call"]); const [defaultOn, setDefaultOn] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState("empty"); const [code, setCode] = useState(CODE_TEMPLATES.empty.code); const [isSaving, setIsSaving] = useState(false); const [isTesting, setIsTesting] = useState(false); const [testExpanded, setTestExpanded] = useState(false); - const [testInput, setTestInput] = useState('{"texts": ["Hello, my SSN is 123-45-6789"], "images": [], "tools": [], "tool_calls": [], "structured_messages": [], "model": "gpt-4"}'); + + // Test input examples for pre_call and post_call + const TEST_INPUT_EXAMPLES = { + pre_call: { + name: "Pre-call (Request)", + data: { + texts: ["Hello, my SSN is 123-45-6789"], + images: [], + tools: [ + { + type: "function", + function: { + name: "get_weather", + description: "Get the current weather in a location", + parameters: { + type: "object", + properties: { + location: { type: "string", description: "City name" } + }, + required: ["location"] + } + } + } + ], + tool_calls: [], + structured_messages: [ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "Hello, my SSN is 123-45-6789" } + ], + model: "gpt-4" + } + }, + post_call: { + name: "Post-call (Response)", + data: { + texts: ["The weather in San Francisco is 72°F and sunny."], + images: [], + tools: [], + tool_calls: [ + { + id: "call_abc123", + type: "function", + function: { + name: "get_weather", + arguments: "{\"location\": \"San Francisco\"}" + } + } + ], + structured_messages: [], + model: "gpt-4" + } + } + }; + + const [testInput, setTestInput] = useState(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2)); const [testResult, setTestResult] = useState(null); const [copiedPrimitive, setCopiedPrimitive] = useState(null); const textareaRef = useRef(null); @@ -145,18 +244,35 @@ const CustomCodeModal: React.FC = ({ setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code); }; - // Reset form when modal opens + // Normalize mode from API (string or string[]) to string[] + const normalizeMode = (m: string | string[] | undefined): string[] => { + if (m === undefined || m === null) return ["pre_call"]; + if (Array.isArray(m)) return m.length ? m : ["pre_call"]; + return [m]; + }; + + // Reset form when modal opens or editData changes useEffect(() => { if (visible) { - setGuardrailName(""); - setMode("pre_call"); - setDefaultOn(false); - setSelectedTemplate("empty"); - setCode(CODE_TEMPLATES.empty.code); + if (editData) { + // Edit mode: populate with existing data + setGuardrailName(editData.guardrail_name || ""); + setMode(normalizeMode(editData.litellm_params?.mode)); + setDefaultOn(editData.litellm_params?.default_on || false); + setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code); + setSelectedTemplate(""); // No template selected in edit mode + } else { + // Create mode: reset to defaults + setGuardrailName(""); + setMode(["pre_call"]); + setDefaultOn(false); + setSelectedTemplate("empty"); + setCode(CODE_TEMPLATES.empty.code); + } setTestResult(null); setTestExpanded(false); } - }, [visible]); + }, [visible, editData]); // Copy primitive to clipboard const copyPrimitive = async (primitive: string) => { @@ -184,7 +300,7 @@ const CustomCodeModal: React.FC = ({ } }; - // Save guardrail + // Save guardrail (create or update) const handleSave = async () => { if (!guardrailName.trim()) { NotificationsManager.fromBackend("Please enter a guardrail name"); @@ -201,25 +317,53 @@ const CustomCodeModal: React.FC = ({ setIsSaving(true); try { - const guardrailData = { - guardrail_name: guardrailName, - litellm_params: { - guardrail: "custom_code", - mode: mode, - default_on: defaultOn, - custom_code: code, - }, - guardrail_info: {}, - }; + if (isEditMode && editData) { + // Update existing guardrail + const updateData: any = { + litellm_params: { + custom_code: code, + }, + }; - await createGuardrailCall(accessToken, guardrailData); - NotificationsManager.success("Custom code guardrail created successfully"); + // Only include changed fields + if (guardrailName !== editData.guardrail_name) { + updateData.guardrail_name = guardrailName; + } + const existingMode = normalizeMode(editData.litellm_params?.mode); + const modeChanged = + mode.length !== existingMode.length || + mode.some((m, i) => m !== existingMode[i]); + if (modeChanged) { + updateData.litellm_params.mode = mode; + } + if (defaultOn !== editData.litellm_params?.default_on) { + updateData.litellm_params.default_on = defaultOn; + } + + await updateGuardrailCall(accessToken, editData.guardrail_id, updateData); + NotificationsManager.success("Custom code guardrail updated successfully"); + } else { + // Create new guardrail + const guardrailData = { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "custom_code", + mode: mode, + default_on: defaultOn, + custom_code: code, + }, + guardrail_info: {}, + }; + + await createGuardrailCall(accessToken, guardrailData); + NotificationsManager.success("Custom code guardrail created successfully"); + } onSuccess(); onClose(); } catch (error) { - console.error("Failed to create guardrail:", error); + console.error("Failed to save guardrail:", error); NotificationsManager.fromBackend( - "Failed to create guardrail: " + (error instanceof Error ? error.message : String(error)) + `Failed to ${isEditMode ? "update" : "create"} guardrail: ` + (error instanceof Error ? error.message : String(error)) ); } finally { setIsSaving(false); @@ -252,10 +396,20 @@ const CustomCodeModal: React.FC = ({ parsedInput.texts = []; } + // Use first request-like or response-like mode for test input_type + const requestModes = ["pre_call", "pre_mcp_call"]; + const responseModes = ["post_call", "post_mcp_call"]; + const testInputType: "request" | "response" = + mode.some((m) => requestModes.includes(m)) + ? "request" + : mode.some((m) => responseModes.includes(m)) + ? "response" + : "request"; + const response = await testCustomCodeGuardrail(accessToken, { custom_code: code, test_input: parsedInput, - input_type: mode as "request" | "response", + input_type: testInputType, request_data: { model: "test-model", metadata: {}, @@ -289,7 +443,7 @@ const CustomCodeModal: React.FC = ({ open={visible} onCancel={onClose} footer={null} - width={1200} + width={1400} className="custom-code-modal" closable={true} destroyOnClose @@ -297,7 +451,9 @@ const CustomCodeModal: React.FC = ({
{/* Header */}
-

Create Custom Guardrail

+

+ {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"} +

Define custom logic using Python-like syntax

@@ -311,14 +467,16 @@ const CustomCodeModal: React.FC = ({ placeholder="e.g., block-pii-custom" />
-
- +
+