From 62d5d96e124afda9edfdd86f4ac71583facfba59 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 12:05:37 +0530 Subject: [PATCH] Fix: skip health check for MCP integration with passthrough token authentication --- .../mcp_server/mcp_server_manager.py | 4 +- .../types/mcp_server/mcp_server_manager.py | 22 ++ .../mcp_server/test_mcp_server_manager.py | 234 ++++++++++++++++++ 3 files changed, 258 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 5c72bfbc13..67a7d4f47f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2463,8 +2463,8 @@ class MCPServerManager: # Check if we should skip health check based on auth configuration should_skip_health_check = False - # Skip if auth_type is oauth2 - if server.needs_user_oauth_token: + # Skip if server requires per-user authentication (OAuth2 or passthrough auth) + if server.requires_per_user_auth: should_skip_health_check = True # Skip if auth_type is not none and authentication_token is missing elif ( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 2cd385c5bf..7f99fd526c 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -65,3 +65,25 @@ class MCPServer(BaseModel): def needs_user_oauth_token(self) -> bool: """True if this is an OAuth2 server that relies on per-user tokens (no client_credentials).""" return self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials + + @property + def requires_per_user_auth(self) -> bool: + """ + True if this server requires per-user/per-request authentication. + This includes: + - OAuth2 servers without client credentials + - Servers with auth_type=none but extra_headers configured for auth passthrough + + Health checks should be skipped for these servers since they cannot + authenticate without user-provided credentials. + """ + # OAuth2 without client credentials + if self.needs_user_oauth_token: + return True + + # PAT passthrough: auth_type is none but extra_headers includes auth headers + if self.auth_type == MCPAuth.none and self.extra_headers: + auth_header_names = {"authorization", "x-api-key", "api-key", "apikey"} + return any(h.lower() in auth_header_names for h in self.extra_headers) + + return False 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 464e523832..c105052479 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 @@ -1036,6 +1036,240 @@ class TestMCPServerManager: assert result.status == "healthy" assert result.health_check_error is None + @pytest.mark.asyncio + async def test_health_check_skips_passthrough_auth_with_authorization_header(self): + """Test that health check is skipped for servers with passthrough Authorization header""" + manager = MCPServerManager() + + # Mock server with auth_type=none and Authorization in extra_headers (passthrough auth) + server = MCPServer( + server_id="github-server", + name="github-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://github-server.com", + extra_headers=["Authorization"], # Passthrough auth configured + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called (health check should be skipped) + manager._create_mcp_client = AsyncMock() + + # Perform health check + result = await manager.health_check_server("github-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "github-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_skips_passthrough_auth_with_api_key_header(self): + """Test that health check is skipped for servers with passthrough x-api-key header""" + manager = MCPServerManager() + + # Mock server with auth_type=none and x-api-key in extra_headers + server = MCPServer( + server_id="sourcegraph-server", + name="sourcegraph-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://sourcegraph-server.com", + extra_headers=["x-api-key"], # Passthrough auth configured + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called + manager._create_mcp_client = AsyncMock() + + # Perform health check + result = await manager.health_check_server("sourcegraph-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "sourcegraph-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_runs_when_no_passthrough_auth(self): + """Test that health check runs normally for servers with auth_type=none but no passthrough headers""" + manager = MCPServerManager() + + # Mock server with auth_type=none but no extra_headers (no passthrough auth) + server = MCPServer( + server_id="public-server", + name="public-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://public-server.com", + extra_headers=None, # No passthrough auth + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = AsyncMock(return_value=mock_client) + + # Perform health check + result = await manager.health_check_server("public-server") + + # Verify that client WAS created (health check should run) + manager._create_mcp_client.assert_called_once() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "public-server" + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_runs_when_extra_headers_no_auth(self): + """Test that health check runs when extra_headers exist but don't include auth headers""" + manager = MCPServerManager() + + # Mock server with extra_headers but no auth-related headers + server = MCPServer( + server_id="custom-server", + name="custom-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://custom-server.com", + extra_headers=["X-Custom-Header", "X-Request-ID"], # Non-auth headers + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = AsyncMock(return_value=mock_client) + + # Perform health check + result = await manager.health_check_server("custom-server") + + # Verify that client WAS created (health check should run) + manager._create_mcp_client.assert_called_once() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "custom-server" + assert result.status == "healthy" + assert result.health_check_error is None + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_oauth2(self): + """Test that requires_per_user_auth returns True for OAuth2 without client credentials""" + # OAuth2 without client credentials + server = MCPServer( + server_id="oauth-server", + name="oauth-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth-server.com", + client_id=None, + client_secret=None, + token_url=None, + ) + assert server.requires_per_user_auth is True + assert server.needs_user_oauth_token is True + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_oauth2_with_client_creds(self): + """Test that requires_per_user_auth returns False for OAuth2 with client credentials""" + # OAuth2 with client credentials + server = MCPServer( + server_id="oauth-server", + name="oauth-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth-server.com", + client_id="client-id", + client_secret="client-secret", + token_url="http://oauth-server.com/token", + ) + assert server.requires_per_user_auth is False + assert server.has_client_credentials is True + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_passthrough_auth(self): + """Test that requires_per_user_auth returns True for passthrough auth (auth_type=none + Authorization header)""" + # Passthrough auth with Authorization header + server = MCPServer( + server_id="github-server", + name="github-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://github-server.com", + extra_headers=["Authorization"], + ) + assert server.requires_per_user_auth is True + + # Passthrough auth with x-api-key header + server2 = MCPServer( + server_id="sourcegraph-server", + name="sourcegraph-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://sourcegraph-server.com", + extra_headers=["x-api-key"], + ) + assert server2.requires_per_user_auth is True + + # Passthrough auth with api-key header (case insensitive) + server3 = MCPServer( + server_id="api-server", + name="api-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://api-server.com", + extra_headers=["API-Key"], + ) + assert server3.requires_per_user_auth is True + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_no_passthrough(self): + """Test that requires_per_user_auth returns False when no passthrough auth is configured""" + # auth_type=none but no extra_headers + server = MCPServer( + server_id="public-server", + name="public-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://public-server.com", + extra_headers=None, + ) + assert server.requires_per_user_auth is False + + # auth_type=none with non-auth extra_headers + server2 = MCPServer( + server_id="custom-server", + name="custom-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://custom-server.com", + extra_headers=["X-Custom-Header", "X-Request-ID"], + ) + assert server2.requires_per_user_auth is False + @pytest.mark.asyncio async def test_register_openapi_tools_includes_static_headers(self, tmp_path): """Ensure OpenAPI-to-MCP tool calls include server.static_headers (Issue #19341)."""