diff --git a/src/fetch/pyproject.toml b/src/fetch/pyproject.toml index e2d0d38d0c..3dbd317f94 100644 --- a/src/fetch/pyproject.toml +++ b/src/fetch/pyproject.toml @@ -38,3 +38,8 @@ dev-dependencies = ["pyright>=1.1.389", "ruff>=0.7.3", "pytest>=8.0.0", "pytest- [tool.pytest.ini_options] testpaths = ["tests"] asyncio_mode = "auto" + +[tool.pyright] +pythonVersion = "3.10" +typeCheckingMode = "basic" +include = ["src", "tests"] diff --git a/src/fetch/src/mcp_server_fetch/server.py b/src/fetch/src/mcp_server_fetch/server.py index b42c7b1f6b..924e05700a 100644 --- a/src/fetch/src/mcp_server_fetch/server.py +++ b/src/fetch/src/mcp_server_fetch/server.py @@ -1,6 +1,8 @@ +from httpx import HTTPError from typing import Annotated, Tuple from urllib.parse import urlparse, urlunparse +import httpx import markdownify import readabilipy.simple_json from mcp.shared.exceptions import McpError @@ -23,6 +25,8 @@ DEFAULT_USER_AGENT_AUTONOMOUS = "ModelContextProtocol/1.0 (Autonomous; +https://github.com/modelcontextprotocol/servers)" DEFAULT_USER_AGENT_MANUAL = "ModelContextProtocol/1.0 (User-Specified; +https://github.com/modelcontextprotocol/servers)" +# Define a reasonable default safety cap (e.g., 2MB) +MAX_RESPONSE_BYTES = 2 * 1024 * 1024 def extract_content_from_html(html: str) -> str: """Extract and convert HTML content to Markdown format. @@ -32,12 +36,16 @@ def extract_content_from_html(html: str) -> str: Returns: Simplified markdown version of the content + + Raises: + ValueError: If the page content cannot be simplified """ ret = readabilipy.simple_json.simple_json_from_html_string( html, use_readability=True ) if not ret["content"]: - return "Page failed to be simplified from HTML" + raise ValueError("Page failed to be simplified from HTML") + content = markdownify.markdownify( ret["content"], heading_style=markdownify.ATX, @@ -68,11 +76,9 @@ async def check_may_autonomously_fetch_url(url: str, user_agent: str, proxy_url: Check if the URL can be fetched by the user agent according to the robots.txt file. Raises a McpError if not. """ - from httpx import AsyncClient, HTTPError - robot_txt_url = get_robots_txt_url(url) - async with AsyncClient(proxy=proxy_url) as client: + async with httpx.AsyncClient(proxy=proxy_url) as client: try: response = await client.get( robot_txt_url, @@ -107,34 +113,46 @@ async def check_may_autonomously_fetch_url(url: str, user_agent: str, proxy_url: f"The assistant can tell the user that they can try manually fetching the page by using the fetch prompt within their UI.", )) - async def fetch_url( url: str, user_agent: str, force_raw: bool = False, proxy_url: str | None = None ) -> Tuple[str, str]: """ Fetch the URL and return the content in a form ready for the LLM, as well as a prefix string with status information. """ - from httpx import AsyncClient, HTTPError - - async with AsyncClient(proxy=proxy_url) as client: + async with httpx.AsyncClient(proxy=proxy_url, follow_redirects=True, timeout=30.0) as client: try: - response = await client.get( - url, - follow_redirects=True, - headers={"User-Agent": user_agent}, - timeout=30, - ) - except HTTPError as e: + async with client.stream("GET", url, headers={"User-Agent": user_agent}) as response: + if response.status_code >= 400: + raise McpError(ErrorData( + code=INTERNAL_ERROR, + message=f"Failed to fetch {url} - status code {response.status_code}", + )) + + content_length = response.headers.get("Content-Length") + if content_length and int(content_length) > MAX_RESPONSE_BYTES: + raise McpError(ErrorData( + code=INVALID_PARAMS, + message=f"Response exceeds maximum allowed size of {MAX_RESPONSE_BYTES} bytes.", + )) + + chunks = [] + bytes_read = 0 + + async for chunk in response.aiter_bytes(): + bytes_read += len(chunk) + if bytes_read > MAX_RESPONSE_BYTES: + raise McpError(ErrorData( + code=INVALID_PARAMS, + message=f"Response body exceeded maximum limit of {MAX_RESPONSE_BYTES} bytes during download.", + )) + chunks.append(chunk) + + page_raw = b"".join(chunks).decode("utf-8", errors="replace") + content_type = response.headers.get("content-type", "") + + except httpx.HTTPError as e: raise McpError(ErrorData(code=INTERNAL_ERROR, message=f"Failed to fetch {url}: {e!r}")) - if response.status_code >= 400: - raise McpError(ErrorData( - code=INTERNAL_ERROR, - message=f"Failed to fetch {url} - status code {response.status_code}", - )) - - page_raw = response.text - content_type = response.headers.get("content-type", "") is_page_html = ( "" in result + with pytest.raises(ValueError, match="Page failed to be simplified from HTML"): + extract_content_from_html(html) class TestCheckMayAutonomouslyFetchUrl: @@ -190,9 +190,7 @@ class TestFetchUrl: @pytest.mark.asyncio async def test_fetch_html_page(self): """Test fetching an HTML page returns markdown content.""" - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.text = """ + html_content = """
@@ -202,20 +200,23 @@ async def test_fetch_html_page(self): """ + mock_response = AsyncMock() + mock_response.status_code = 200 mock_response.headers = {"content-type": "text/html"} + mock_response.aiter_bytes = MagicMock() + mock_response.aiter_bytes.return_value.__aiter__.return_value = [html_content.encode("utf-8")] with patch("httpx.AsyncClient") as mock_client_class: mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_client_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_class.return_value.__aexit__ = AsyncMock(return_value=None) + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client content, prefix = await fetch_url( "https://example.com/page", DEFAULT_USER_AGENT_AUTONOMOUS ) - # HTML is processed, so we check it returns something assert isinstance(content, str) assert prefix == "" @@ -223,16 +224,17 @@ async def test_fetch_html_page(self): async def test_fetch_html_page_raw(self): """Test fetching an HTML page with raw=True returns original HTML.""" html_content = "

Test

" - mock_response = MagicMock() + mock_response = AsyncMock() mock_response.status_code = 200 - mock_response.text = html_content mock_response.headers = {"content-type": "text/html"} + mock_response.aiter_bytes = MagicMock() + mock_response.aiter_bytes.return_value.__aiter__.return_value = [html_content.encode("utf-8")] with patch("httpx.AsyncClient") as mock_client_class: mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_client_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_class.return_value.__aexit__ = AsyncMock(return_value=None) + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client content, prefix = await fetch_url( "https://example.com/page", @@ -247,16 +249,17 @@ async def test_fetch_html_page_raw(self): async def test_fetch_json_returns_raw(self): """Test fetching JSON content returns raw content.""" json_content = '{"key": "value"}' - mock_response = MagicMock() + mock_response = AsyncMock() mock_response.status_code = 200 - mock_response.text = json_content mock_response.headers = {"content-type": "application/json"} + mock_response.aiter_bytes = MagicMock() + mock_response.aiter_bytes.return_value.__aiter__.return_value = [json_content.encode("utf-8")] with patch("httpx.AsyncClient") as mock_client_class: mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_client_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_class.return_value.__aexit__ = AsyncMock(return_value=None) + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client content, prefix = await fetch_url( "https://api.example.com/data", @@ -269,14 +272,15 @@ async def test_fetch_json_returns_raw(self): @pytest.mark.asyncio async def test_fetch_404_raises_error(self): """Test that 404 response raises McpError.""" - mock_response = MagicMock() + mock_response = AsyncMock() mock_response.status_code = 404 + mock_response.headers = {} with patch("httpx.AsyncClient") as mock_client_class: mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_client_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_class.return_value.__aexit__ = AsyncMock(return_value=None) + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client with pytest.raises(McpError): await fetch_url( @@ -287,14 +291,15 @@ async def test_fetch_404_raises_error(self): @pytest.mark.asyncio async def test_fetch_500_raises_error(self): """Test that 500 response raises McpError.""" - mock_response = MagicMock() + mock_response = AsyncMock() mock_response.status_code = 500 + mock_response.headers = {} with patch("httpx.AsyncClient") as mock_client_class: mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_client_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_class.return_value.__aexit__ = AsyncMock(return_value=None) + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client with pytest.raises(McpError): await fetch_url( @@ -305,16 +310,17 @@ async def test_fetch_500_raises_error(self): @pytest.mark.asyncio async def test_fetch_with_proxy(self): """Test that proxy URL is passed to client.""" - mock_response = MagicMock() + mock_response = AsyncMock() mock_response.status_code = 200 - mock_response.text = '{"data": "test"}' mock_response.headers = {"content-type": "application/json"} + mock_response.aiter_bytes = MagicMock() + mock_response.aiter_bytes.return_value.__aiter__.return_value = [b'{"data": "test"}'] with patch("httpx.AsyncClient") as mock_client_class: mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_client_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_class.return_value.__aexit__ = AsyncMock(return_value=None) + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client await fetch_url( "https://example.com/data", @@ -323,4 +329,49 @@ async def test_fetch_with_proxy(self): ) # Verify AsyncClient was called with proxy - mock_client_class.assert_called_once_with(proxy="http://proxy.example.com:8080") + mock_client_class.assert_called_once_with( + proxy="http://proxy.example.com:8080", + follow_redirects=True, + timeout=30.0 + ) + + @pytest.mark.asyncio + async def test_fetch_exceeds_content_length_limit(self): + """Test that exceeding Content-Length limit raises McpError.""" + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Length": str(10 * 1024 * 1024)} # 10MB + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client + + with pytest.raises(McpError, match="Response exceeds maximum allowed size"): + await fetch_url( + "https://example.com/largefile", + DEFAULT_USER_AGENT_AUTONOMOUS + ) + + @pytest.mark.asyncio + async def test_fetch_exceeds_download_limit(self): + """Test that exceeding download limit during streaming raises McpError.""" + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.headers = {} # No content-length + mock_response.aiter_bytes = MagicMock() + # Return chunks that exceed 2MB total + mock_response.aiter_bytes.return_value.__aiter__.return_value = [b"a" * (1024 * 1024)] * 3 + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.stream = MagicMock() + mock_client.stream.return_value.__aenter__ = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client + + with pytest.raises(McpError, match="Response body exceeded maximum limit"): + await fetch_url( + "https://example.com/largefile", + DEFAULT_USER_AGENT_AUTONOMOUS + )