From 36999b23ee976726631035054ca2f7df3196c62a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 4 Mar 2026 13:07:25 +0530 Subject: [PATCH 01/63] [Chore] update mcp documentation for header forwarding --- docs/my-website/docs/mcp.md | 57 +++++++++++++++++++ docs/my-website/docs/mcp_control.md | 8 +-- .../src/components/mcp_tools/mcp_connect.tsx | 2 +- 3 files changed, 62 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index fcbb31c07d3..c7789201579 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -870,6 +870,63 @@ asyncio.run(main()) [Learn more about customer management →](./proxy/customers) +## Calling the Proxy's /v1/responses Endpoint + +When calling your LiteLLM Proxy's `/v1/responses` endpoint to use MCP tools, **always use `server_url: "litellm_proxy"`** in the tools array. This tells the proxy to use its configured MCP servers. + +:::important Do not use the full proxy URL +Using `server_url: "https://your-proxy.com/mcp"` is incorrect when the request is already going to the proxy. The proxy needs the literal value `litellm_proxy` to route to its configured MCP servers. +::: + +```bash title="Correct: Using litellm_proxy" showLineNumbers +curl --location 'https://your-proxy.com/v1/responses' \ +--header 'Content-Type: application/json' \ +--header "Authorization: Bearer $LITELLM_API_KEY" \ +--data '{ + "model": "gpt-4", + "tools": [ + { + "type": "mcp", + "server_label": "litellm", + "server_url": "litellm_proxy", + "require_approval": "never" + } + ], + "input": "Run available tools", + "tool_choice": "required" +}' +``` + +### Sending Custom Headers to MCP Servers + +To pass custom headers (e.g., API keys, auth tokens) to specific MCP servers, use either: + +**Option 1: Request headers** – Add `x-mcp-{server_alias}-{header_name}` to your request headers. The proxy forwards these to the matching MCP server. + +```bash +# Send Authorization header to the "weather2" MCP server +--header 'x-mcp-weather2-authorization: Bearer your-token' + +# Send custom header to the "github" MCP server +--header 'x-mcp-github-x-api-key: your-api-key' +``` + +**Option 2: Headers in tool config** – Include a `headers` object in the tool definition. These are merged with request headers. + +```json +{ + "type": "mcp", + "server_label": "litellm", + "server_url": "litellm_proxy", + "require_approval": "never", + "headers": { + "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", + "x-mcp-servers": "Zapier_MCP,dev-group", + "x-mcp-weather2-authorization": "Bearer your-weather-api-token" + } +} +``` + ## Using your MCP with client side credentials Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP. diff --git a/docs/my-website/docs/mcp_control.md b/docs/my-website/docs/mcp_control.md index 96c71ef9278..ccaa37f9497 100644 --- a/docs/my-website/docs/mcp_control.md +++ b/docs/my-website/docs/mcp_control.md @@ -323,7 +323,7 @@ curl --location '/v1/responses' \ { "type": "mcp", "server_label": "litellm", - "server_url": "/dev_group/mcp", + "server_url": "litellm_proxy", "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY" @@ -335,7 +335,7 @@ curl --location '/v1/responses' \ }' ``` -This example uses URL namespacing to access all servers in the "dev_group" access group. +This example uses the `x-mcp-servers` header to access all servers in the "dev_group" access group. Use `server_url: "litellm_proxy"` when calling the proxy's `/v1/responses` endpoint—do not use the full proxy URL. @@ -423,7 +423,7 @@ curl --location '/v1/responses' \ { "type": "mcp", "server_label": "litellm", - "server_url": "/mcp/", + "server_url": "litellm_proxy", "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", @@ -436,7 +436,7 @@ curl --location '/v1/responses' \ }' ``` -This configuration restricts the request to only use tools from the specified MCP servers. +This configuration restricts the request to only use tools from the specified MCP servers. Use `server_url: "litellm_proxy"` when calling the proxy's `/v1/responses` endpoint. diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx index c48b9a755b7..1c82859e062 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx @@ -256,7 +256,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] { "type": "mcp", "server_label": "litellm", - "server_url": "${proxyBaseUrl}/mcp", + "server_url": "litellm_proxy", "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_VIRTUAL_KEY", From 96b75be03d1db6e4957183061fb20e97163318ee Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Mar 2026 20:13:14 -0800 Subject: [PATCH 02/63] [Feature] RBAC for Vector Stores and Agents Add proxy-admin-configurable toggles to restrict internal users (and optionally team admins) from accessing agent and vector store management features. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/agent_endpoints/endpoints.py | 20 ++- litellm/proxy/common_utils/rbac_utils.py | 126 ++++++++++++++ .../proxy_setting_endpoints.py | 60 +++++-- .../management_endpoints.py | 11 ++ .../proxy/agent_endpoints/test_agent_rbac.py | 84 ++++++++++ .../proxy/common_utils/test_rbac_utils.py | 156 ++++++++++++++++++ .../test_vector_store_rbac.py | 121 ++++++++++++++ .../components/SidebarProvider.tsx | 12 ++ .../AdminSettings/UISettings/UISettings.tsx | 136 +++++++++++++++ .../src/components/leftnav.tsx | 8 +- 10 files changed, 720 insertions(+), 14 deletions(-) create mode 100644 litellm/proxy/common_utils/rbac_utils.py create mode 100644 tests/litellm/proxy/agent_endpoints/test_agent_rbac.py create mode 100644 tests/litellm/proxy/common_utils/test_rbac_utils.py create mode 100644 tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 65674d01be7..80c55f634f7 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -16,6 +16,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.types.agents import ( AgentConfig, @@ -69,6 +70,8 @@ async def get_agents( Returns: List[AgentResponse] """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, @@ -179,6 +182,8 @@ async def create_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -233,7 +238,10 @@ async def create_agent( dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) -async def get_agent_by_id(agent_id: str): +async def get_agent_by_id( + agent_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Get a specific agent by ID @@ -243,6 +251,8 @@ async def get_agent_by_id(agent_id: str): -H "Authorization: Bearer " ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -319,6 +329,8 @@ async def update_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -410,6 +422,8 @@ async def patch_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -484,6 +498,8 @@ async def delete_agent( } ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -763,6 +779,8 @@ async def get_agent_daily_activity( """ Get daily activity for specific agents or all accessible agents. """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: diff --git a/litellm/proxy/common_utils/rbac_utils.py b/litellm/proxy/common_utils/rbac_utils.py new file mode 100644 index 00000000000..2b187d18065 --- /dev/null +++ b/litellm/proxy/common_utils/rbac_utils.py @@ -0,0 +1,126 @@ +""" +RBAC utility helpers for feature-level access control. + +These helpers are used by agent and vector store endpoints to enforce +proxy-admin-configurable toggles that restrict access for internal users. +""" + +from typing import TYPE_CHECKING + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth + +if TYPE_CHECKING: + pass + + +def _is_user_team_admin_for_any_team( + user_api_key_dict: UserAPIKeyAuth, + teams: list, +) -> bool: + """ + Return True if the user is an admin member in at least one of the given teams. + + Args: + user_api_key_dict: The authenticated user. + teams: List of Prisma team records (from litellm_teamtable.find_many). + """ + for team in teams: + team_obj = LiteLLM_TeamTable(**team.model_dump()) + for member in team_obj.members_with_roles: + if ( + member.user_id is not None + and member.user_id == user_api_key_dict.user_id + and member.role == "admin" + ): + return True + return False + + +async def check_feature_access_for_user( + user_api_key_dict: UserAPIKeyAuth, + feature_name: str, +) -> None: + """ + Raise HTTP 403 if the user's role is blocked from accessing the given feature + by the UI settings stored in general_settings. + + Args: + user_api_key_dict: The authenticated user. + feature_name: Either "agents" or "vector_stores". + """ + # Proxy admins (and view-only admins) are never blocked. + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ): + return + + from litellm.proxy.proxy_server import general_settings + + disable_flag = f"disable_{feature_name}_for_internal_users" + allow_team_admins_flag = f"allow_{feature_name}_for_team_admins" + + if not general_settings.get(disable_flag, False): + # Feature is not disabled — allow all authenticated users. + return + + # Feature is disabled. Check if team admins are exempted. + if general_settings.get(allow_team_admins_flag, False): + is_team_admin = await _check_if_team_admin(user_api_key_dict) + if is_team_admin: + return + + raise HTTPException( + status_code=403, + detail={ + "error": f"Access to {feature_name} is disabled for your role. Contact your proxy admin." + }, + ) + + +async def _check_if_team_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Return True if the user is a team admin in any team. + Mirrors the logic in management_endpoints/common_utils._user_has_admin_privileges + but scoped to team-admin check only. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None or user_api_key_dict.user_id is None: + return False + + from litellm.caching import DualCache + from litellm.proxy.auth.auth_checks import get_user_object + + try: + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + user_id_upsert=False, + proxy_logging_obj=None, + ) + + if user_obj is None: + return False + + if user_obj.teams is None or len(user_obj.teams) == 0: + return False + + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": user_obj.teams}} + ) + + return _is_user_team_admin_for_any_team(user_api_key_dict, teams) + + except Exception as e: + verbose_proxy_logger.debug( + f"rbac_utils: error checking team admin status for user " + f"{user_api_key_dict.user_id}: {e}" + ) + return False diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index ceda08d520a..8991dc5fd5c 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -104,6 +104,26 @@ class UISettings(BaseModel): description="If enabled, shows the Projects feature in the UI sidebar and the project field in key management.", ) + disable_agents_for_internal_users: bool = Field( + default=False, + description="If true, internal users cannot access agent management endpoints or the Agents page in the UI.", + ) + + allow_agents_for_team_admins: bool = Field( + default=False, + description="If true, team admins are exempt from the agents disable restriction (only takes effect when disable_agents_for_internal_users is true).", + ) + + disable_vector_stores_for_internal_users: bool = Field( + default=False, + description="If true, internal users cannot access vector store management endpoints or the Vector Stores page in the UI.", + ) + + allow_vector_stores_for_team_admins: bool = Field( + default=False, + description="If true, team admins are exempt from the vector stores disable restriction (only takes effect when disable_vector_stores_for_internal_users is true).", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -119,6 +139,10 @@ ALLOWED_UI_SETTINGS_FIELDS = { "require_auth_for_public_ai_hub", "forward_client_headers_to_llm_api", "enable_projects_ui", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", } @@ -976,14 +1000,20 @@ async def get_ui_settings(): k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS } - # Sync forward_client_headers_to_llm_api into general_settings so the proxy - # picks it up at runtime (covers server restart scenarios). - if "forward_client_headers_to_llm_api" in ui_settings: + # Sync runtime flags into general_settings so the proxy picks them up + # at runtime (covers server restart scenarios). + _runtime_flags = [ + "forward_client_headers_to_llm_api", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", + ] + _flags_to_sync = {k: ui_settings[k] for k in _runtime_flags if k in ui_settings} + if _flags_to_sync: from litellm.proxy.proxy_server import general_settings - general_settings["forward_client_headers_to_llm_api"] = ui_settings[ - "forward_client_headers_to_llm_api" - ] + general_settings.update(_flags_to_sync) # Build config-like object for schema helper config: Dict[str, Any] = {"litellm_settings": {"ui_settings": ui_settings}} @@ -1048,14 +1078,20 @@ async def update_ui_settings( }, ) - # Sync forward_client_headers_to_llm_api to general_settings so the proxy - # picks it up at runtime (general_settings is checked in pre-call utils). - if "forward_client_headers_to_llm_api" in ui_settings: + # Sync runtime flags to general_settings so the proxy picks them up + # at runtime (general_settings is checked in pre-call utils). + _runtime_flags = [ + "forward_client_headers_to_llm_api", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", + ] + _flags_to_sync = {k: ui_settings[k] for k in _runtime_flags if k in ui_settings} + if _flags_to_sync: from litellm.proxy.proxy_server import general_settings - general_settings["forward_client_headers_to_llm_api"] = ui_settings[ - "forward_client_headers_to_llm_api" - ] + general_settings.update(_flags_to_sync) return { "message": "UI settings updated successfully", diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cccbb51f47b..068f4217e0f 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, @@ -439,6 +440,8 @@ async def new_vector_store( - vector_store_description: Optional[str] - Description of the vector store - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client try: @@ -506,6 +509,8 @@ async def list_vector_stores( - page: int - Page number for pagination (default: 1) - page_size: int - Number of items per page (default: 100) """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client vector_store_map: Dict[str, LiteLLM_ManagedVectorStore] = {} @@ -605,6 +610,8 @@ async def delete_vector_store( Parameters: - vector_store_id: str - ID of the vector store to delete """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -687,6 +694,8 @@ async def get_vector_store_info( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return a single vector store's details""" + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -770,6 +779,8 @@ async def update_vector_store( Update vector store details in both database and in-memory registry. The updated data is immediately synchronized to the in-memory registry. """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client from litellm.types.router import GenericLiteLLMParams diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py new file mode 100644 index 00000000000..a863201ddb5 --- /dev/null +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -0,0 +1,84 @@ +""" +Tests for RBAC enforcement on agent endpoints. + +Verifies that check_feature_access_for_user is called and that a 403 is +raised when agents are disabled for internal users. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id=user_id, + ) + + +def _make_admin_user(user_id: str = "admin-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id=user_id, + ) + + +# --------------------------------------------------------------------------- +# get_agents +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_get_agents_blocked_for_internal_user_when_disabled(): + """get_agents should raise 403 when agents are disabled for internal users.""" + from litellm.proxy.agent_endpoints.endpoints import get_agents + + user = _make_internal_user() + gs = {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False} + + request_mock = MagicMock() + with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True): + with pytest.raises(HTTPException) as exc_info: + await get_agents(request=request_mock, user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_agents_allowed_when_not_disabled(): + """get_agents should not raise RBAC 403 when agents are not disabled.""" + from litellm.proxy.agent_endpoints.endpoints import get_agents + + user = _make_internal_user() + request_mock = MagicMock() + + with patch.dict("litellm.proxy.proxy_server.general_settings", {}, clear=True): + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + MagicMock(get_agent_list=MagicMock(return_value=[])), + ): + with patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + new=AsyncMock(return_value=[]), + ): + result = await get_agents(request=request_mock, user_api_key_dict=user) + assert result == [] + + +# --------------------------------------------------------------------------- +# get_agent_daily_activity +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_get_agent_daily_activity_blocked_when_disabled(): + from litellm.proxy.agent_endpoints.endpoints import get_agent_daily_activity + + user = _make_internal_user() + gs = {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False} + + with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True): + with pytest.raises(HTTPException) as exc_info: + await get_agent_daily_activity(user_api_key_dict=user) + assert exc_info.value.status_code == 403 diff --git a/tests/litellm/proxy/common_utils/test_rbac_utils.py b/tests/litellm/proxy/common_utils/test_rbac_utils.py new file mode 100644 index 00000000000..997a2e19b77 --- /dev/null +++ b/tests/litellm/proxy/common_utils/test_rbac_utils.py @@ -0,0 +1,156 @@ +""" +Tests for litellm/proxy/common_utils/rbac_utils.py + +Covers check_feature_access_for_user for agents and vector_stores features. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user + + +def _make_user(role: str, user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=role, user_id=user_id) + + +# general_settings is imported from litellm.proxy.proxy_server inside the +# function, so we patch it via patch.dict on the original dict. +_GS_PATH = "litellm.proxy.proxy_server.general_settings" + + +# --------------------------------------------------------------------------- +# Proxy admin is always allowed +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_proxy_admin_always_allowed(): + user = _make_user(LitellmUserRoles.PROXY_ADMIN.value) + with patch.dict(_GS_PATH, {"disable_agents_for_internal_users": True}): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_proxy_admin_view_only_always_allowed(): + user = _make_user(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) + with patch.dict(_GS_PATH, {"disable_agents_for_internal_users": True}): + await check_feature_access_for_user(user, "agents") + + +# --------------------------------------------------------------------------- +# Feature not disabled — everyone allowed +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_feature_not_disabled_allows_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict(_GS_PATH, {}, clear=True): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_feature_not_disabled_allows_vector_stores(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict(_GS_PATH, {"disable_vector_stores_for_internal_users": False}, clear=True): + await check_feature_access_for_user(user, "vector_stores") + + +# --------------------------------------------------------------------------- +# Feature disabled, team-admin exemption OFF — internal user blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_agents_disabled_blocks_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False}, + clear=True, + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "agents") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_blocks_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": False}, + clear=True, + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "vector_stores") + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# Feature disabled, allow_team_admins ON — team admin allowed, non-admin blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_agents_disabled_team_admin_allowed(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="team-admin-user") + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=True), + ): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_agents_disabled_non_team_admin_blocked(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="regular-user") + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=False), + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "agents") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_team_admin_allowed(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="team-admin-user") + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=True), + ): + await check_feature_access_for_user(user, "vector_stores") + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_non_team_admin_blocked(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="regular-user") + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=False), + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "vector_stores") + assert exc_info.value.status_code == 403 diff --git a/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py b/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py new file mode 100644 index 00000000000..3eb49bdf114 --- /dev/null +++ b/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py @@ -0,0 +1,121 @@ +""" +Tests for RBAC enforcement on vector store management endpoints. + +Verifies that check_feature_access_for_user is called and that a 403 is +raised when vector stores are disabled for internal users. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id=user_id, + ) + + +_DISABLED_GS = { + "disable_vector_stores_for_internal_users": True, + "allow_vector_stores_for_team_admins": False, +} + +_ENABLED_GS: dict = {} + + +# --------------------------------------------------------------------------- +# list_vector_stores +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_list_vector_stores_blocked_when_disabled(): + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + user = _make_internal_user() + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with pytest.raises(HTTPException) as exc_info: + await list_vector_stores(user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_list_vector_stores_allowed_when_not_disabled(): + """list_vector_stores should not raise 403 when vector stores are not disabled.""" + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + import litellm + user = _make_internal_user() + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[]) + + raised_403 = False + with patch.dict("litellm.proxy.proxy_server.general_settings", _ENABLED_GS, clear=True): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(litellm, "vector_store_registry", None): + with patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[]), + ): + try: + await list_vector_stores(user_api_key_dict=user) + except HTTPException as e: + if e.status_code == 403: + raised_403 = True + assert not raised_403, "Should not raise 403 when vector stores are not disabled" + + +# --------------------------------------------------------------------------- +# new_vector_store +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_new_vector_store_blocked_when_disabled(): + from litellm.proxy.vector_store_endpoints.management_endpoints import new_vector_store + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + + user = _make_internal_user() + vs = LiteLLM_ManagedVectorStore(vector_store_id="vs-1", custom_llm_provider="openai") # type: ignore[call-arg] + + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with pytest.raises(HTTPException) as exc_info: + await new_vector_store(vector_store=vs, user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# Admin user is never blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_list_vector_stores_admin_not_blocked(): + """Proxy admin should never be blocked, even when vector stores are disabled.""" + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + import litellm + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id="admin-1", + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[]) + + raised_403 = False + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(litellm, "vector_store_registry", None): + with patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[]), + ): + try: + await list_vector_stores(user_api_key_dict=admin) + except HTTPException as e: + if e.status_code == 403: + raised_403 = True + assert not raised_403, "Admin should not be blocked even when vector stores are disabled" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx index 17f62a20f7d..7dcc3fa8a1a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx @@ -15,6 +15,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side const { accessToken } = useAuthorized(); const [enabledPagesInternalUsers, setEnabledPagesInternalUsers] = useState(null); const [enableProjectsUI, setEnableProjectsUI] = useState(false); + const [disableAgentsForInternalUsers, setDisableAgentsForInternalUsers] = useState(false); + const [disableVectorStoresForInternalUsers, setDisableVectorStoresForInternalUsers] = useState(false); useEffect(() => { const fetchUISettings = async () => { @@ -39,6 +41,14 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side if (settings?.values?.enable_projects_ui !== undefined) { setEnableProjectsUI(Boolean(settings.values.enable_projects_ui)); } + + if (settings?.values?.disable_agents_for_internal_users !== undefined) { + setDisableAgentsForInternalUsers(Boolean(settings.values.disable_agents_for_internal_users)); + } + + if (settings?.values?.disable_vector_stores_for_internal_users !== undefined) { + setDisableVectorStoresForInternalUsers(Boolean(settings.values.disable_vector_stores_for_internal_users)); + } } catch (error) { console.error("[SidebarProvider] Failed to fetch UI settings:", error); } @@ -54,6 +64,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side collapsed={sidebarCollapsed} enabledPagesInternalUsers={enabledPagesInternalUsers} enableProjectsUI={enableProjectsUI} + disableAgentsForInternalUsers={disableAgentsForInternalUsers} + disableVectorStoresForInternalUsers={disableVectorStoresForInternalUsers} /> ); }; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 5d99dd2969d..dfc66d3484d 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -19,9 +19,15 @@ export default function UISettings() { const forwardClientHeadersProperty = schema?.properties?.forward_client_headers_to_llm_api; const enableProjectsUIProperty = schema?.properties?.enable_projects_ui; const enabledPagesProperty = schema?.properties?.enabled_ui_pages_internal_users; + const disableAgentsProperty = schema?.properties?.disable_agents_for_internal_users; + const allowAgentsTeamAdminsProperty = schema?.properties?.allow_agents_for_team_admins; + const disableVectorStoresProperty = schema?.properties?.disable_vector_stores_for_internal_users; + const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); + const isAgentsDisabled = Boolean(values.disable_agents_for_internal_users); + const isVectorStoresDisabled = Boolean(values.disable_vector_stores_for_internal_users); const handleToggle = (checked: boolean) => { updateSettings( @@ -105,6 +111,62 @@ export default function UISettings() { ); }; + const handleToggleDisableAgents = (checked: boolean) => { + updateSettings( + { disable_agents_for_internal_users: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleAllowAgentsTeamAdmins = (checked: boolean) => { + updateSettings( + { allow_agents_for_team_admins: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleDisableVectorStores = (checked: boolean) => { + updateSettings( + { disable_vector_stores_for_internal_users: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleAllowVectorStoresTeamAdmins = (checked: boolean) => { + updateSettings( + { allow_vector_stores_for_team_admins: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + return ( {isLoading ? ( @@ -211,6 +273,80 @@ export default function UISettings() { + {/* Agents access control */} + + + + Disable agents for internal users + {disableAgentsProperty?.description && ( + {disableAgentsProperty.description} + )} + + + + + + + + Allow agents for team admins + + {allowAgentsTeamAdminsProperty?.description && ( + {allowAgentsTeamAdminsProperty.description} + )} + + + + + + {/* Vector Stores access control */} + + + + Disable vector stores for internal users + {disableVectorStoresProperty?.description && ( + {disableVectorStoresProperty.description} + )} + + + + + + + + Allow vector stores for team admins + + {allowVectorStoresTeamAdminsProperty?.description && ( + {allowVectorStoresTeamAdminsProperty.description} + )} + + + + + {/* Page Visibility for Internal Users */} = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI }) => { +const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI, disableAgentsForInternalUsers, disableVectorStoresForInternalUsers }) => { const { userId, accessToken, userRole } = useAuthorized(); const { data: organizations } = useOrganizations(); @@ -450,6 +452,10 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse // Hide Projects page if enableProjectsUI is not enabled if (item.key === "projects" && !enableProjectsUI) return false; + // Hide agents and vector-stores pages for non-admin users when disabled + if (!isAdmin && item.key === "agents" && disableAgentsForInternalUsers) return false; + if (!isAdmin && item.key === "vector-stores" && disableVectorStoresForInternalUsers) return false; + // Existing role check if (item.roles && !item.roles.includes(userRole)) return false; From 028e6871dd5f8611f84c1e2dc853f44e506e5a92 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:27:51 +0530 Subject: [PATCH 03/63] feat(agents): add static_headers and extra_headers fields to schema and types Add two new fields to LiteLLM_AgentsTable: - static_headers (Json): admin-configured headers always sent to the backend agent - extra_headers (String[]): header names to extract from the client request and forward Extend AgentConfig, PatchAgentRequest, and AgentResponse with the same fields. Also remove duplicate spec_path field from LiteLLM_MCPServerTable. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/schema.prisma | 3 ++- litellm/types/agents.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 43972724ecc..6f4ef0c24b6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -63,6 +63,8 @@ model LiteLLM_AgentsTable { agent_name String @unique litellm_params Json? agent_card_params Json + static_headers Json? @default("{}") + extra_headers String[] @default([]) agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) @@ -305,7 +307,6 @@ model LiteLLM_MCPServerTable { registration_url String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) - spec_path String? is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 3ad898b1935..7879cae9ff6 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -179,6 +179,8 @@ class AgentConfig(TypedDict, total=False): agent_card_params: Required[AgentCard] litellm_params: Dict[str, Any] # allow for any future litellm params object_permission: AgentObjectPermission + static_headers: Optional[Dict[str, str]] + extra_headers: Optional[List[str]] class PatchAgentRequest(TypedDict, total=False): @@ -186,6 +188,8 @@ class PatchAgentRequest(TypedDict, total=False): agent_card_params: AgentCard litellm_params: Dict[str, Any] object_permission: AgentObjectPermission + static_headers: Optional[Dict[str, str]] + extra_headers: Optional[List[str]] # Request/Response models for CRUD endpoints @@ -197,6 +201,8 @@ class AgentResponse(BaseModel): litellm_params: Optional[Dict[str, Any]] = None agent_card_params: Dict[str, Any] object_permission: Optional[Dict[str, Any]] = None + static_headers: Optional[Dict[str, str]] = None + extra_headers: Optional[List[str]] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None created_by: Optional[str] = None From 07ee1e9886f54b773f4d9de7e3c8181e90d30d6e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:01 +0530 Subject: [PATCH 04/63] feat(agents): persist static_headers and extra_headers in agent registry Update add_agent_to_db, patch_agent_in_db, and update_agent_in_db to read and write the two new header fields when creating or updating agents. Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/agent_endpoints/agent_registry.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 159c9fb93d9..550182f966f 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -128,6 +128,14 @@ class AgentRegistry: agent_copy, None, prisma_client ) + # Serialize static_headers + static_headers_obj = agent.get("static_headers") + static_headers_val: Optional[str] = ( + safe_dumps(dict(static_headers_obj)) if static_headers_obj else None + ) + + extra_headers_val: Optional[List[str]] = agent.get("extra_headers") + create_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, @@ -137,6 +145,10 @@ class AgentRegistry: "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), } + if static_headers_val is not None: + create_data["static_headers"] = static_headers_val + if extra_headers_val is not None: + create_data["extra_headers"] = extra_headers_val if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id @@ -214,6 +226,12 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) + if agent.get("static_headers") is not None: + update_data["static_headers"] = safe_dumps( + dict(agent.get("static_headers")) # type: ignore + ) + if agent.get("extra_headers") is not None: + update_data["extra_headers"] = agent.get("extra_headers") if agent.get("object_permission") is not None: agent_copy = dict(augment_agent) existing_object_permission_id = existing_agent.get( @@ -281,6 +299,15 @@ class AgentRegistry: ) agent_card_params: str = safe_dumps(agent_card_params_dict) + # Serialize static_headers for update + static_headers_obj_u = agent.get("static_headers") + static_headers_val_u: Optional[str] = ( + safe_dumps(dict(static_headers_obj_u)) + if static_headers_obj_u is not None + else None + ) + extra_headers_val_u: Optional[List[str]] = agent.get("extra_headers") + update_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, @@ -288,6 +315,10 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } + if static_headers_val_u is not None: + update_data["static_headers"] = static_headers_val_u + if extra_headers_val_u is not None: + update_data["extra_headers"] = extra_headers_val_u if agent.get("object_permission") is not None: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} From 16a30b55f5493bbf0754aac0dc4ea4c54b681804 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:11 +0530 Subject: [PATCH 05/63] feat(agents): add merge_agent_headers utility Mirrors merge_mcp_headers from the MCP server utils. Dynamic headers come first; static (admin-configured) headers overlay and win on conflict. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/agent_endpoints/utils.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 litellm/proxy/agent_endpoints/utils.py diff --git a/litellm/proxy/agent_endpoints/utils.py b/litellm/proxy/agent_endpoints/utils.py new file mode 100644 index 00000000000..2b968de54be --- /dev/null +++ b/litellm/proxy/agent_endpoints/utils.py @@ -0,0 +1,27 @@ +"""Utility helpers for A2A agent endpoints.""" + +from typing import Dict, Mapping, Optional + + +def merge_agent_headers( + *, + dynamic_headers: Optional[Mapping[str, str]] = None, + static_headers: Optional[Mapping[str, str]] = None, +) -> Optional[Dict[str, str]]: + """Merge outbound HTTP headers for A2A agent calls. + + Merge rules: + - Start with ``dynamic_headers`` (values extracted from the incoming client request). + - Overlay ``static_headers`` (admin-configured per agent). + + If both contain the same key, ``static_headers`` wins. + """ + merged: Dict[str, str] = {} + + if dynamic_headers: + merged.update({str(k): str(v) for k, v in dynamic_headers.items()}) + + if static_headers: + merged.update({str(k): str(v) for k, v in static_headers.items()}) + + return merged or None From 20a4eea27e71cfc5933670b73747fb46d66dd41d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:28 +0530 Subject: [PATCH 06/63] feat(agents): forward custom headers to backend A2A agents In invoke_agent_a2a: - Extract admin-configured extra_headers from client request by name - Extract convention-based headers (x-a2a-{agent_id/name}-{header}) from client request - Merge with static_headers (static wins on conflict) - Pass merged headers down to asend_message and _handle_stream_message In asend_message / asend_message_streaming: - Accept agent_extra_headers kwarg - Overlay onto LiteLLM internal headers before creating the httpx client Co-Authored-By: Claude Sonnet 4.6 --- litellm/a2a_protocol/main.py | 14 ++++++- .../proxy/agent_endpoints/a2a_endpoints.py | 37 ++++++++++++++++++- 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 485b57e311b..6ac88d3a430 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -169,6 +169,7 @@ async def asend_message( api_base: Optional[str] = None, litellm_params: Optional[Dict[str, Any]] = None, agent_id: Optional[str] = None, + agent_extra_headers: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> LiteLLMSendMessageResponse: """ @@ -250,9 +251,12 @@ async def asend_message( "Either a2a_client or api_base is required for standard A2A flow" ) trace_id = trace_id or str(uuid.uuid4()) - extra_headers = {"X-LiteLLM-Trace-Id": trace_id} + extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: extra_headers["X-LiteLLM-Agent-Id"] = agent_id + # Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones) + if agent_extra_headers: + extra_headers.update(agent_extra_headers) a2a_client = await create_a2a_client( base_url=api_base, extra_headers=extra_headers ) @@ -426,6 +430,7 @@ async def asend_message_streaming( agent_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, proxy_server_request: Optional[Dict[str, Any]] = None, + agent_extra_headers: Optional[Dict[str, str]] = None, ) -> AsyncIterator[Any]: """ Async: Send a streaming message to an A2A agent. @@ -507,7 +512,12 @@ async def asend_message_streaming( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - a2a_client = await create_a2a_client(base_url=api_base) + streaming_extra_headers: Optional[Dict[str, str]] = None + if agent_extra_headers: + streaming_extra_headers = dict(agent_extra_headers) + a2a_client = await create_a2a_client( + base_url=api_base, extra_headers=streaming_extra_headers + ) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 6bcee14f29e..344070d17fc 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -6,13 +6,14 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from typing import Any, Optional +from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.utils import all_litellm_params @@ -55,6 +56,7 @@ async def _handle_stream_message( metadata: Optional[dict] = None, proxy_server_request: Optional[dict] = None, *, + agent_extra_headers: Optional[Dict[str, str]] = None, user_api_key_dict: Optional[UserAPIKeyAuth] = None, request_data: Optional[dict] = None, proxy_logging_obj: Optional[Any] = None, @@ -105,6 +107,7 @@ async def _handle_stream_message( agent_id=agent_id, metadata=metadata, proxy_server_request=proxy_server_request, + agent_extra_headers=agent_extra_headers, ) if ( @@ -385,6 +388,36 @@ async def invoke_agent_a2a( version=version, ) + # Build merged headers for the backend agent + static_headers: Dict[str, str] = dict(agent.static_headers or {}) + + raw_headers = dict(request.headers) + normalized = {k.lower(): v for k, v in raw_headers.items()} + + dynamic_headers: Dict[str, str] = {} + + # 1. Admin-configured extra_headers: forward named headers from client request + if agent.extra_headers: + for header_name in agent.extra_headers: + val = normalized.get(header_name.lower()) + if val is not None: + dynamic_headers[header_name] = val + + # 2. Convention-based forwarding: x-a2a-{agent_id_or_name}-{header_name} + # Matches both agent_id (UUID) and agent_name (alias), case-insensitive. + for alias in (agent.agent_id.lower(), agent.agent_name.lower()): + prefix = f"x-a2a-{alias}-" + for key, val in normalized.items(): + if key.startswith(prefix): + header_name = key[len(prefix) :] + if header_name: + dynamic_headers[header_name] = val + + agent_extra_headers = merge_agent_headers( + dynamic_headers=dynamic_headers or None, + static_headers=static_headers or None, + ) + # Route through SDK functions if method == "message/send": from a2a.types import MessageSendParams, SendMessageRequest @@ -401,6 +434,7 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), litellm_logging_obj=logging_obj, + agent_extra_headers=agent_extra_headers, ) response = await proxy_logging_obj.post_call_success_hook( @@ -425,6 +459,7 @@ async def invoke_agent_a2a( agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), + agent_extra_headers=agent_extra_headers, user_api_key_dict=user_api_key_dict, request_data=data, proxy_logging_obj=proxy_logging_obj, From 6e9c7c4a8dd8ddce1b911d77e2009aac3de5f9d3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:36 +0530 Subject: [PATCH 07/63] feat(agents): add Prisma migration for agent header columns ALTER TABLE LiteLLM_AgentsTable to add: - static_headers JSONB DEFAULT '{}' - extra_headers TEXT[] DEFAULT ARRAY[]::TEXT[] Co-Authored-By: Claude Sonnet 4.6 --- .../20260305000000_add_agent_headers/migration.sql | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql new file mode 100644 index 00000000000..acb35baba96 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql @@ -0,0 +1,5 @@ +-- Add static_headers and extra_headers to LiteLLM_AgentsTable + +ALTER TABLE "LiteLLM_AgentsTable" + ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}', + ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; From fd53678898b71f6da4384e5984a4d1308f2ee060 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:48 +0530 Subject: [PATCH 08/63] test(agents): add tests for A2A custom header forwarding Covers: - Static headers forwarded to backend - Dynamic headers extracted by name (extra_headers config) - Convention-based x-a2a-{agent_id/name}-{header} forwarding - Static headers win over dynamic on conflict - Unrelated x-a2a- prefixes are not forwarded - No-header case leaves existing behaviour unchanged - merge_agent_headers utility unit tests Co-Authored-By: Claude Sonnet 4.6 --- .../agent_endpoints/test_agent_headers.py | 339 ++++++++++++++++++ 1 file changed, 339 insertions(+) create mode 100644 tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py new file mode 100644 index 00000000000..b52c0afb0c0 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py @@ -0,0 +1,339 @@ +""" +Unit tests for A2A agent custom header forwarding. + +Tests cover: +- Static headers forwarded to backend agent +- Dynamic headers extracted from client request and forwarded +- Static headers win over dynamic on conflict +- No headers configured — existing behavior unchanged +- merge_agent_headers utility +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +# --------------------------------------------------------------------------- +# Helper: build a minimal mock agent +# --------------------------------------------------------------------------- + +def _make_mock_agent( + static_headers=None, + extra_headers=None, + url="http://backend-agent:10001", +): + mock_agent = MagicMock() + mock_agent.agent_id = "agent-123" + mock_agent.agent_card_params = {"url": url, "name": "Test Agent"} + mock_agent.litellm_params = {} + mock_agent.static_headers = static_headers or {} + mock_agent.extra_headers = extra_headers or [] + return mock_agent + + +def _make_mock_request(extra_headers=None, method="message/send"): + """Build a mock FastAPI Request with configurable headers.""" + mock_request = MagicMock() + headers = {"content-type": "application/json"} + if extra_headers: + headers.update(extra_headers) + mock_request.headers = headers + mock_request.json = AsyncMock( + return_value={ + "jsonrpc": "2.0", + "id": "test-id", + "method": method, + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } + }, + } + ) + return mock_request + + +def _make_a2a_types_module(): + """Return (module, MessageSendParams, SendMessageRequest, SendStreamingMessageRequest).""" + try: + from a2a.types import ( + MessageSendParams, + SendMessageRequest, + SendStreamingMessageRequest, + ) + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = MessageSendParams + mock_a2a_types.SendMessageRequest = SendMessageRequest + mock_a2a_types.SendStreamingMessageRequest = SendStreamingMessageRequest + return mock_a2a_types + except ImportError: + pass + + def _make_cls(name): + class MockCls: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + self._kwargs = kwargs + + def model_dump(self, mode="json", exclude_none=False): + result = dict(self._kwargs) + if exclude_none: + result = {k: v for k, v in result.items() if v is not None} + return result + + MockCls.__name__ = name + return MockCls + + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = _make_cls("MessageSendParams") + mock_a2a_types.SendMessageRequest = _make_cls("SendMessageRequest") + mock_a2a_types.SendStreamingMessageRequest = _make_cls( + "SendStreamingMessageRequest" + ) + return mock_a2a_types + + +async def _invoke(mock_agent, mock_request, mock_asend_message): + """Run invoke_agent_a2a with standard patches applied.""" + from litellm.proxy._types import UserAPIKeyAuth + + mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + mock_fastapi_response = MagicMock() + mock_a2a_types = _make_a2a_types_module() + + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "jsonrpc": "2.0", + "id": "test-id", + "result": {"status": "success"}, + } + + with patch( + "litellm.proxy.agent_endpoints.a2a_endpoints._get_agent", + return_value=mock_agent, + ), patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + new_callable=AsyncMock, + return_value=True, + ), patch( + "litellm.proxy.common_request_processing.add_litellm_data_to_request", + side_effect=lambda data, **kw: data, + ), patch( + "litellm.a2a_protocol.asend_message", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_asend, patch( + "litellm.a2a_protocol.create_a2a_client", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ), patch( + "litellm.proxy.proxy_server.proxy_config", + MagicMock(), + ), patch( + "litellm.proxy.proxy_server.version", + "1.0.0", + ), patch.dict( + sys.modules, + {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, + ), patch( + "litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", + True, + ): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + return mock_asend + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_static_headers_forwarded(): + """Static headers configured on the agent are passed to asend_message.""" + mock_agent = _make_mock_agent( + static_headers={"Authorization": "Bearer token123"} + ) + mock_request = _make_mock_request() + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None, "agent_extra_headers should not be None" + assert headers.get("Authorization") == "Bearer token123" + + +@pytest.mark.asyncio +async def test_dynamic_headers_forwarded(): + """Dynamic headers listed in extra_headers are extracted from the client request.""" + mock_agent = _make_mock_agent(extra_headers=["x-api-key"]) + mock_request = _make_mock_request(extra_headers={"x-api-key": "secret"}) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("x-api-key") == "secret" + + +@pytest.mark.asyncio +async def test_static_overrides_dynamic(): + """When the same header appears in both static and dynamic, static wins.""" + mock_agent = _make_mock_agent( + static_headers={"Authorization": "Bearer static-token"}, + extra_headers=["Authorization"], + ) + # Client sends a different value for Authorization + mock_request = _make_mock_request( + extra_headers={"Authorization": "Bearer dynamic-token"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("Authorization") == "Bearer static-token" + + +@pytest.mark.asyncio +async def test_no_headers(): + """When no headers are configured, agent_extra_headers is None and behaviour is unchanged.""" + mock_agent = _make_mock_agent() # no static_headers or extra_headers + mock_request = _make_mock_request() + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is None + + +# --------------------------------------------------------------------------- +# Convention-based x-a2a-{agent_id/name}-{header_name} tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_convention_header_by_agent_name(): + """x-a2a-{agent_name}-{header} is forwarded using the agent name alias.""" + mock_agent = _make_mock_agent() + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-my-agent-authorization": "Bearer conv-token"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("authorization") == "Bearer conv-token" + + +@pytest.mark.asyncio +async def test_convention_header_by_agent_id(): + """x-a2a-{agent_id}-{header} is forwarded using the agent UUID.""" + mock_agent = _make_mock_agent() + mock_agent.agent_id = "abc-123" + mock_agent.agent_name = "other-name" + mock_request = _make_mock_request( + extra_headers={"x-a2a-abc-123-x-api-key": "id-secret"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("x-api-key") == "id-secret" + + +@pytest.mark.asyncio +async def test_convention_header_static_still_wins(): + """Static headers still override convention-based dynamic headers.""" + mock_agent = _make_mock_agent( + static_headers={"authorization": "Bearer static-wins"} + ) + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-my-agent-authorization": "Bearer conv-value"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("authorization") == "Bearer static-wins" + + +@pytest.mark.asyncio +async def test_convention_unrelated_prefix_not_forwarded(): + """Headers that start with x-a2a- but target a different agent are ignored.""" + mock_agent = _make_mock_agent() + mock_agent.agent_id = "agent-abc" + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-other-agent-authorization": "Bearer wrong"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is None + + +# --------------------------------------------------------------------------- +# Direct unit test for the merge utility +# --------------------------------------------------------------------------- + + +def test_merge_agent_headers_util_dynamic_only(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(dynamic_headers={"x-key": "val"}) + assert result == {"x-key": "val"} + + +def test_merge_agent_headers_util_static_only(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(static_headers={"Authorization": "Bearer tok"}) + assert result == {"Authorization": "Bearer tok"} + + +def test_merge_agent_headers_util_static_wins(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers( + dynamic_headers={"Authorization": "dynamic", "x-extra": "d"}, + static_headers={"Authorization": "static"}, + ) + assert result == {"Authorization": "static", "x-extra": "d"} + + +def test_merge_agent_headers_util_none_returns_none(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers() + assert result is None + + +def test_merge_agent_headers_util_empty_dicts_returns_none(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(dynamic_headers={}, static_headers={}) + assert result is None From 36d279ab42c20185d435d502f18f895487249ab3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:34:11 +0530 Subject: [PATCH 09/63] feat(ui/agents): add Authentication Headers section to agent create/edit form Add a new "Authentication Headers" panel to AgentFormFields: - Static Headers: key-value Form.List (always sent to the backend agent, static wins on conflict with dynamic) - Forward Client Headers: Select[tags] of header names to extract from the client request and forward (extra_headers) Update buildAgentDataFromForm to serialize both fields for the API. Update parseAgentForForm to deserialize them back for editing. Covers both the create wizard (add_agent_form) and the edit view (agent_info). Co-Authored-By: Claude Sonnet 4.6 --- .../src/components/agents/agent_config.ts | 26 +++++++ .../components/agents/agent_form_fields.tsx | 70 ++++++++++++++++++- 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index f85c4daac66..01041c5cee4 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -269,6 +269,23 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { agentData.litellm_params = params; } + // static_headers: convert [{header, value}, ...] → {header: value, ...} + if (Array.isArray(values.static_headers) && values.static_headers.length > 0) { + const staticHeaders: Record = {}; + values.static_headers.forEach((entry: { header?: string; value?: string }) => { + const key = entry?.header?.trim(); + if (key) staticHeaders[key] = entry?.value ?? ""; + }); + if (Object.keys(staticHeaders).length > 0) { + agentData.static_headers = staticHeaders; + } + } + + // extra_headers: already an array of strings from Select tags + if (Array.isArray(values.extra_headers) && values.extra_headers.length > 0) { + agentData.extra_headers = values.extra_headers; + } + return agentData; }; @@ -302,5 +319,14 @@ export const parseAgentForForm = (agent: any) => { cost_per_query: agent.litellm_params?.cost_per_query, input_cost_per_token: agent.litellm_params?.input_cost_per_token, output_cost_per_token: agent.litellm_params?.output_cost_per_token, + // static_headers: {key: value} → [{header, value}, ...] + static_headers: agent.static_headers + ? Object.entries(agent.static_headers as Record).map(([header, value]) => ({ + header, + value, + })) + : [], + // extra_headers: already an array of strings + extra_headers: agent.extra_headers ?? [], }; }; diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index d5429d2a3b5..42e55b8c56f 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -1,7 +1,7 @@ import React from "react"; -import { Form, Input, Switch, Collapse } from "antd"; +import { Form, Input, Switch, Collapse, Select, Space, Tooltip } from "antd"; import { Button as AntButton } from "antd"; -import { PlusOutlined, MinusCircleOutlined } from "@ant-design/icons"; +import { PlusOutlined, MinusCircleOutlined, InfoCircleOutlined } from "@ant-design/icons"; import { AGENT_FORM_CONFIG, SKILL_FIELD_CONFIG } from "./agent_config"; import CostConfigFields from "./cost_config_fields"; @@ -188,6 +188,72 @@ const AgentFormFields: React.FC = ({ showAgentName = true, ))} )} + + {/* Authentication Headers */} + {shouldShow("auth_headers") && ( + + {/* Static Headers */} + + Static Headers{" "} + + + + + } + > + + {(fields, { add, remove }) => ( + <> + {fields.map(({ key, name, ...restField }) => ( + + + + + + + + remove(name)} style={{ color: "#ff4d4f" }} /> + + ))} + add()} icon={} style={{ width: "100%" }}> + Add Static Header + + + )} + + + + {/* Extra Headers (dynamic forwarding) */} + + Forward Client Headers{" "} + + + + + } + name="extra_headers" + > +