feat(fetch): implement streaming byte cap to prevent resource exhaustion

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

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 "<error>Page failed to be simplified from HTML</error>"
+        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 = (
         "<html" in page_raw[:100] or "text/html" in content_type or not content_type
     )

src/fetch/tests/test_server.py

@@ -81,11 +81,11 @@ def test_html_with_links(self):
         result = extract_content_from_html(html)
         assert "Example" in result
 
-    def test_empty_content_returns_error(self):
-        """Test that empty/invalid HTML returns error message."""
+    def test_empty_content_raises_error(self):
+        """Test that empty/invalid HTML raises ValueError."""
         html = ""
-        result = extract_content_from_html(html)
-        assert "<error>" 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 = """
         <html>
         <body>
             <article>
@@ -202,37 +200,41 @@ async def test_fetch_html_page(self):
         </body>
         </html>
         """
+        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 == ""
 
     @pytest.mark.asyncio
     async def test_fetch_html_page_raw(self):
         """Test fetching an HTML page with raw=True returns original HTML."""
         html_content = "<html><body><h1>Test</h1></body></html>"
-        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
+                )