From c8ef01bfbeddf29b871cfc65efa42b38d8433c26 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 16:43:53 -0700 Subject: [PATCH] feat(mcp): make allowed_response_headers configurable from the database and dashboard The field was config.yaml only, so a server created through the dashboard or the REST API could never surface upstream response headers. Adds the column to all three prisma schemas plus a migration, carries it through the create, edit and temporary-record paths, the row -> MCPServer build, and the dashboard's edit form and detail view. The unsupported-transport warning now also fires for servers loaded from the database, since the dashboard can create the same misconfiguration as config.yaml. The column is a Prisma String[], which rejects null, and the edit form clears a field to null, so an explicit null normalizes to [] alongside allowed_tools. --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/mcp_server.py | 1 + litellm/proxy/_experimental/mcp_server/db.py | 5 +- .../mcp_server/mcp_server_manager.py | 2 + litellm/proxy/_types.py | 2 + .../mcp_management_endpoints.py | 3 ++ litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + .../mcp_server/test_db_credentials.py | 38 ++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 51 +++++++++++++++++++ .../test_mcp_management_endpoints.py | 27 ++++++++++ .../_components/MCPPermissionManagement.tsx | 30 +++++++++++ .../_components/mcp_server_edit.tsx | 2 + .../_components/mcp_server_view.tsx | 10 ++++ .../src/components/mcp_tools/types.tsx | 1 + 16 files changed, 175 insertions(+), 2 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_allowed_response_headers_to_mcp_server_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_allowed_response_headers_to_mcp_server_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_allowed_response_headers_to_mcp_server_table/migration.sql new file mode 100644 index 00000000000..d54e3eb355f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_allowed_response_headers_to_mcp_server_table/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "allowed_response_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f842bf13da9..813fab487fa 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -311,6 +311,7 @@ model LiteLLM_MCPServerTable { tool_name_to_display_name Json? @default("{}") tool_name_to_description Json? @default("{}") extra_headers String[] @default([]) + allowed_response_headers String[] @default([]) static_headers Json? @default("{}") // Admin-configured environment variables interpolated into static_headers // via ${NAME} syntax. Stored as an array of diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 23b26bd8e89..f52e452c4e5 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -67,6 +67,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): tool_name_to_display_name: Optional[Dict[str, str]] = None tool_name_to_description: Optional[Dict[str, str]] = None extra_headers: List[str] = Field(default_factory=list) + allowed_response_headers: list[str] = Field(default_factory=list) mcp_info: Optional[MCPInfo] = None static_headers: Optional[Dict[str, str]] = None env_vars: Optional[List[MCPEnvVar]] = None diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index d55eb3ac014..56d3dc76b26 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -289,8 +289,9 @@ def _prepare_mcp_server_data( data_dict.pop("alias", None) # Prisma ``allowed_tools`` is a required String[]; ``null`` is invalid. # The UI sends null to clear a whitelist — treat that as ``[]``. - if "allowed_tools" in data_dict and data_dict["allowed_tools"] is None: - data_dict["allowed_tools"] = [] + for list_field in ("allowed_tools", "allowed_response_headers"): + if list_field in data_dict and data_dict[list_field] is None: + data_dict[list_field] = [] # Json map fields use ``@default("{}")``; explicit null means clear overrides. for json_map_field in ( "tool_name_to_display_name", diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2a33dd37909..677c08a87d9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1768,6 +1768,7 @@ class MCPServerManager: authentication_token=auth_value, mcp_info=mcp_info, extra_headers=getattr(mcp_server, "extra_headers", None), + allowed_response_headers=getattr(mcp_server, "allowed_response_headers", None), static_headers=static_headers_dict, env_vars=env_vars_list, client_id=client_id_value or getattr(mcp_server, "client_id", None), @@ -1825,6 +1826,7 @@ class MCPServerManager: max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None), ) _warn_internal_delegate_pkce_if_applicable(new_server, source="database") + _warn_response_headers_unsupported_transport_if_applicable(new_server, source="database") if persist_discovered_endpoints: await self._persist_discovered_obo_token_url( server_id=mcp_server.server_id, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 053f78a3698..ff73a30d100 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1256,6 +1256,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): tool_name_to_display_name: Optional[Dict[str, str]] = None tool_name_to_description: Optional[Dict[str, str]] = None extra_headers: Optional[List[str]] = None + allowed_response_headers: list[str] | None = None static_headers: Optional[Dict[str, str]] = None env_vars: Optional[List[MCPEnvVar]] = None instructions: Optional[str] = None @@ -1362,6 +1363,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): tool_name_to_display_name: Optional[Dict[str, str]] = None tool_name_to_description: Optional[Dict[str, str]] = None extra_headers: Optional[List[str]] = None + allowed_response_headers: list[str] | None = None static_headers: Optional[Dict[str, str]] = None env_vars: Optional[List[MCPEnvVar]] = None instructions: Optional[str] = None diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d920ee474cc..afa4b60222d 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -533,6 +533,7 @@ if MCP_AVAILABLE: sanitized.spec_path = None sanitized.static_headers = None sanitized.extra_headers = [] + sanitized.allowed_response_headers = [] sanitized.env = {} sanitized.command = None sanitized.args = [] @@ -577,6 +578,7 @@ if MCP_AVAILABLE: sanitized.command = None sanitized.args = [] sanitized.extra_headers = [] + sanitized.allowed_response_headers = [] sanitized.allowed_tools = [] sanitized.mcp_access_groups = [] sanitized.teams = [] @@ -683,6 +685,7 @@ if MCP_AVAILABLE: mcp_access_groups=payload.mcp_access_groups, allowed_tools=payload.allowed_tools or [], extra_headers=payload.extra_headers or [], + allowed_response_headers=payload.allowed_response_headers or [], mcp_info=payload.mcp_info, static_headers=payload.static_headers, command=payload.command, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f842bf13da9..813fab487fa 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -311,6 +311,7 @@ model LiteLLM_MCPServerTable { tool_name_to_display_name Json? @default("{}") tool_name_to_description Json? @default("{}") extra_headers String[] @default([]) + allowed_response_headers String[] @default([]) static_headers Json? @default("{}") // Admin-configured environment variables interpolated into static_headers // via ${NAME} syntax. Stored as an array of diff --git a/schema.prisma b/schema.prisma index f842bf13da9..813fab487fa 100644 --- a/schema.prisma +++ b/schema.prisma @@ -311,6 +311,7 @@ model LiteLLM_MCPServerTable { tool_name_to_display_name Json? @default("{}") tool_name_to_description Json? @default("{}") extra_headers String[] @default([]) + allowed_response_headers String[] @default([]) static_headers Json? @default("{}") // Admin-configured environment variables interpolated into static_headers // via ${NAME} syntax. Stored as an array of diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 7269774442b..f9ad734ff24 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -958,6 +958,44 @@ def test_prepare_mcp_server_data_create_carries_token_exchange_columns(): assert data["token_exchange_profile"] == "entra_obo" +def test_prepare_mcp_server_data_create_carries_allowed_response_headers(): + """The create path must emit allowed_response_headers as a column value, or a server registered + through the REST API / dashboard could never surface upstream response headers.""" + request = NewMCPServerRequest( + server_name="hdr_write", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + + data = _prepare_mcp_server_data(request) + + assert data["allowed_response_headers"] == ["X-Example-Header"] + + +def test_prepare_mcp_server_data_update_maps_cleared_allowed_response_headers_to_empty_list(): + """The column is a Prisma String[], which rejects null. The edit form clears the field to null, + so an explicit null must become [] rather than reaching the DB and failing the update.""" + request = UpdateMCPServerRequest( + server_id="hdr-1", + allowed_response_headers=None, + ) + + data = _prepare_mcp_server_data(request, exclude_unset=True) + + assert data["allowed_response_headers"] == [] + + +def test_prepare_mcp_server_data_update_omits_allowed_response_headers_when_untouched(): + """A partial update that never mentions the field must not write it, so an edit to an unrelated + field cannot silently wipe a configured allowlist.""" + request = UpdateMCPServerRequest(server_id="hdr-1", alias="renamed") + + data = _prepare_mcp_server_data(request, exclude_unset=True) + + assert "allowed_response_headers" not in data + + def test_prepare_mcp_server_data_update_carries_token_exchange_columns(): """The partial-update path (PUT /v1/mcp/server, exclude_unset) must carry the three token-exchange columns when the caller provides them.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 02cf611b8b9..3eb99250bda 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -6560,6 +6560,57 @@ class TestWarnResponseHeadersUnsupportedTransport: assert not any("allowed_response_headers" in m for m in caplog.messages) + @pytest.mark.asyncio + async def test_loading_such_a_server_from_the_database_emits_the_warning(self, caplog): + """The field is DB-backed, so the dashboard can create the same misconfiguration as config.yaml.""" + manager = MCPServerManager() + table_record = LiteLLM_MCPServerTable( + server_id="hdr-sse-db", + server_name="hdr_sse_db", + url="https://example.com/sse", + transport=MCPTransport.sse, + allowed_response_headers=["X-Example-Header"], + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.build_mcp_server_from_table(table_record) + + assert any("allowed_response_headers is set but the transport is" in m for m in caplog.messages) + + +class TestAllowedResponseHeadersFromDatabase: + """The DB row is the dashboard's storage, so the field must survive the row -> MCPServer build.""" + + @pytest.mark.asyncio + async def test_build_mcp_server_from_table_carries_allowed_response_headers(self): + """Without this the dashboard could save the allowlist and the gateway would silently ignore it.""" + manager = MCPServerManager() + table_record = LiteLLM_MCPServerTable( + server_id="hdr-db-1", + server_name="hdr_db", + url="https://example.com/mcp", + transport=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + + mcp_server = await manager.build_mcp_server_from_table(table_record) + + assert mcp_server.allowed_response_headers == ["X-Example-Header"] + + @pytest.mark.asyncio + async def test_build_mcp_server_from_table_defaults_to_no_headers(self): + manager = MCPServerManager() + table_record = LiteLLM_MCPServerTable( + server_id="hdr-db-2", + server_name="hdr_db_none", + url="https://example.com/mcp", + transport=MCPTransport.http, + ) + + mcp_server = await manager.build_mcp_server_from_table(table_record) + + assert not mcp_server.allowed_response_headers + class TestHasClientCredentialsOAuth2Flow: """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index a669a277d2b..8d2abb8c210 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -5376,3 +5376,30 @@ async def test_edit_mcp_server_snapshot_failure_skips_purge_but_edit_succeeds(): assert result.server_id == server_id mock_purge.assert_not_awaited() + + +def test_temporary_mcp_server_record_carries_allowed_response_headers(): + """The pre-save 'test this server' flow builds an unpersisted record; dropping the field there + would make a temporary server behave differently from the one the admin is about to save.""" + payload = NewMCPServerRequest( + server_name="hdr_tmp", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + + record = mgmt_endpoints._build_temporary_mcp_server_record(payload, created_by="tester") + + assert record.allowed_response_headers == ["X-Example-Header"] + + +def test_temporary_mcp_server_record_defaults_allowed_response_headers_to_empty(): + payload = NewMCPServerRequest( + server_name="hdr_tmp_none", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + ) + + record = mgmt_endpoints._build_temporary_mcp_server_record(payload, created_by="tester") + + assert record.allowed_response_headers == [] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx index aae13d4b467..41662ded325 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx @@ -279,6 +279,36 @@ const MCPPermissionManagement: React.FC = ({ /> + + Allowed Response Headers + + + + {mcpServer?.allowed_response_headers && mcpServer.allowed_response_headers.length > 0 && ( + + {mcpServer.allowed_response_headers.length} configured + + )} + + } + name="allowed_response_headers" + > +