ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: update tests + logic for passing multiple headers

Browse Source
This commit is contained in:
Krrish Dholakia 2025-09-28 09:09:11 -07:00
parent 9403efa511
commit 53d0cbb1b7
2 changed files with 31 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View File

@ -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"

48
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View File

@ -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