From 53d0cbb1b7cb340ae1b549b43387ec9e1f40ac19 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 28 Sep 2025 09:09:11 -0700 Subject: [PATCH] fix: update tests + logic for passing multiple headers --- .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../auth/test_user_api_key_auth_mcp.py | 48 ++++++++++++------- 2 files changed, 31 insertions(+), 19 deletions(-) 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 4b2c3385bb..581da3cec2 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 @@ -186,7 +186,7 @@ class MCPRequestHandler: if len(parts) == 2: server_alias, auth_header_name = parts - # Convert header name to proper case (e.g., "authorization" -> "Authorization") + # Convert common header names to proper case if auth_header_name == "authorization": auth_header_name = "Authorization" 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 7bdefe84c1..bf051085cc 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 @@ -286,7 +286,10 @@ class TestMCPRequestHandler: ], "test-api-key-123", None, - {"github": "Bearer github-token", "zapier_x_api": "zapier-api-key"}, + { + "github": {"Authorization": "Bearer github-token"}, + "zapier_x_api": {"key": "zapier-api-key"}, + }, ), # Test case 10: Both legacy and server-specific auth headers ( @@ -297,7 +300,7 @@ class TestMCPRequestHandler: ], "test-api-key-123", "legacy-token", - {"github": "Bearer github-token"}, + {"github": {"Authorization": "Bearer github-token"}}, ), # Test case 11: Server-specific auth headers with different header types ( @@ -308,7 +311,10 @@ class TestMCPRequestHandler: ], "test-api-key-123", None, - {"deepwiki": "Basic base64-encoded", "custom_x_custom": "custom-value"}, + { + "deepwiki": {"Authorization": "Basic base64-encoded"}, + "custom_x_custom": {"header": "custom-value"}, + }, ), # Test case 12: Case insensitive server-specific headers ( @@ -318,7 +324,7 @@ class TestMCPRequestHandler: ], "test-api-key-123", None, - {"github": "Bearer github-token"}, + {"github": {"Authorization": "Bearer github-token"}}, ), ], ) @@ -365,6 +371,7 @@ class TestMCPRequestHandler: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -751,7 +758,7 @@ class TestMCPCustomHeaderName: } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": "Bearer github-token"} + assert result == {"github": {"Authorization": "Bearer github-token"}} # Test case 3: Multiple server-specific headers headers = Headers( @@ -764,9 +771,9 @@ class TestMCPCustomHeaderName: ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) expected = { - "github": "Bearer github-token", - "zapier_x_api": "zapier-api-key", - "deepwiki": "Basic base64-encoded", + "github": {"Authorization": "Bearer github-token"}, + "zapier_x_api": {"key": "zapier-api-key"}, + "deepwiki": {"Authorization": "Basic base64-encoded"}, } assert result == expected @@ -775,11 +782,14 @@ class TestMCPCustomHeaderName: { "x-litellm-api-key": "test-key", "X-MCP-GITHUB-AUTHORIZATION": "Bearer github-token", - "x-mcp-ZAPIER_x_api-key": "zapier-api-key", + "x-mcp-ZAPIER-x-api-key": "zapier-api-key", } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - expected = {"github": "Bearer github-token", "zapier_x_api": "zapier-api-key"} + expected = { + "github": {"Authorization": "Bearer github-token"}, + "zapier": {"x-api-key": "zapier-api-key"}, + } assert result == expected # Test case 5: Invalid format headers (should be ignored) @@ -792,7 +802,7 @@ class TestMCPCustomHeaderName: } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": "Bearer github-token"} + assert result == {"github": {"Authorization": "Bearer github-token"}} # Test case 6: Edge case - header with multiple hyphens in server alias headers = Headers( @@ -804,8 +814,8 @@ class TestMCPCustomHeaderName: ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) expected = { - "github_mcp": "Bearer github-mcp-token", - "gh_mcp2": "Bearer gh-mcp2-token", + "github_mcp": {"Authorization": "Bearer github-mcp-token"}, + "gh_mcp2": {"Authorization": "Bearer gh-mcp2-token"}, } assert result == expected @@ -817,14 +827,14 @@ class TestMCPCustomHeaderName: } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github_mcp": "Bearer github-mcp-token"} + assert result == {"github_mcp": {"Authorization": "Bearer github-mcp-token"}} # Test case 8: Edge case - empty header value headers = Headers( {"x-litellm-api-key": "test-key", "x-mcp-github-authorization": ""} ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": ""} + assert result == {"github": {"Authorization": ""}} # Test case 9: Edge case - very long header value long_token = "Bearer " + "x" * 1000 @@ -832,7 +842,7 @@ class TestMCPCustomHeaderName: {"x-litellm-api-key": "test-key", "x-mcp-github-authorization": long_token} ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": long_token} + assert result == {"github": {"Authorization": long_token}} # Test case 10: Edge case - special characters in server alias headers = Headers( @@ -844,8 +854,8 @@ class TestMCPCustomHeaderName: ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) expected = { - "github-123": "Bearer github-123-token", - "github_test": "Bearer github-test-token", + "github": {"123-authorization": "Bearer github-123-token"}, + "github_test": {"Authorization": "Bearer github-test-token"}, } assert result == expected @@ -890,6 +900,7 @@ class TestMCPAccessGroupsE2E: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -940,6 +951,7 @@ class TestMCPAccessGroupsE2E: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results