feat: add new mcp tool to screenshot pages

Implements #244

Author: D4Vinci

scrapling/core/ai.py

@@ -3,7 +3,8 @@
 from datetime import datetime, timezone
 from dataclasses import dataclass, field
 
-from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp import FastMCP, Image
+from mcp.types import ImageContent, TextContent
 from pydantic import BaseModel, Field
 
 from scrapling.core.shell import Convertor
@@ -31,6 +32,7 @@
 )
 
 SessionType = Literal["dynamic", "stealthy"]
+ScreenshotType = Literal["png", "jpeg"]
 
 
 class ResponseModel(BaseModel):
@@ -106,14 +108,14 @@ class ScraplingMCPServer:
     def __init__(self):
         self._sessions: Dict[str, _SessionEntry] = {}
 
-    def _get_session(self, session_id: str, expected_type: SessionType) -> _SessionEntry:
-        """Look up a session by ID and validate its type."""
+    def _get_session(self, session_id: str, expected_type: Optional[SessionType]) -> _SessionEntry:
+        """Look up a session by ID, optionally validating its type. Pass `None` to skip the type check."""
         entry = self._sessions.get(session_id)
         if entry is None:
             raise ValueError(f"Session '{session_id}' not found. Use list_sessions to see active sessions.")
         if not entry.session._is_alive:
             raise ValueError(f"Session '{session_id}' is no longer alive. Open a new session.")
-        if entry.session_type != expected_type:
+        if expected_type is not None and entry.session_type != expected_type:
             raise ValueError(
                 f"Session '{session_id}' is a '{entry.session_type}' session, but this tool requires a "
                 f"'{expected_type}' session. Use the matching fetch tool for your session type."
@@ -260,6 +262,69 @@ async def list_sessions(self) -> List[SessionInfo]:
             for sid, entry in self._sessions.items()
         ]
 
+    async def screenshot(
+        self,
+        url: str,
+        session_id: str,
+        image_type: ScreenshotType = "png",
+        full_page: bool = False,
+        quality: Optional[int] = None,
+        wait: int | float = 0,
+        wait_selector: Optional[str] = None,
+        wait_selector_state: SelectorWaitStates = "attached",
+        network_idle: bool = False,
+        timeout: int | float = 30000,
+    ) -> List[ImageContent | TextContent]:
+        """Capture a screenshot of a web page using an existing browser session and return it as an image.
+        A browser session must be opened first with `open_session` (either `dynamic` or `stealthy`); the session ID is then passed here.
+
+        :param url: The URL to navigate to and capture.
+        :param session_id: ID of an open browser session created with `open_session`.
+        :param image_type: Image format. Defaults to "png". Use "jpeg" for smaller file sizes.
+        :param full_page: When True, captures the full scrollable page instead of just the viewport. Defaults to False.
+        :param quality: Image quality (0-100) for JPEG only. Raises if passed with `image_type="png"`.
+        :param wait: Time in milliseconds to wait after page load before capturing. Defaults to 0.
+        :param wait_selector: Optional CSS selector to wait for before capturing.
+        :param wait_selector_state: State to wait for the selector. Defaults to "attached".
+        :param network_idle: Wait for the page until there are no network connections for at least 500 ms.
+        :param timeout: Timeout in milliseconds for page operations. Defaults to 30,000.
+        """
+        if quality is not None and image_type != "jpeg":
+            raise ValueError("'quality' is only valid when 'image_type' is 'jpeg'.")
+
+        entry = self._get_session(session_id, expected_type=None)
+
+        screenshot_kwargs: Dict[str, Any] = {"type": image_type, "full_page": full_page}
+        if quality is not None:
+            screenshot_kwargs["quality"] = quality
+
+        captured: Dict[str, Any] = {}
+
+        async def _capture(page: Any) -> None:
+            try:
+                captured["bytes"] = await page.screenshot(**screenshot_kwargs)
+                captured["url"] = page.url
+            except Exception as exc:
+                captured["error"] = exc
+
+        await entry.session.fetch(
+            url,
+            wait=wait,
+            timeout=timeout,
+            network_idle=network_idle,
+            wait_selector=wait_selector,
+            wait_selector_state=wait_selector_state,
+            page_action=_capture,
+        )
+
+        if "error" in captured:
+            raise captured["error"]
+        if "bytes" not in captured:
+            raise RuntimeError(f"Failed to capture screenshot for {url}")
+
+        image = Image(data=captured["bytes"], format=image_type).to_image_content()
+        return [image, TextContent(type="text", text=captured["url"])]
+
     @staticmethod
     async def get(
         url: str,
@@ -298,7 +363,8 @@ async def get(
         :param headers: Headers to include in the request.
         :param cookies: Cookies to use in the request.
         :param timeout: Number of seconds to wait before timing out.
-        :param follow_redirects: Whether to follow redirects. Defaults to "safe", which follows redirects but rejects those targeting internal/private IPs (SSRF protection). Pass True to follow all redirects without restriction.
+        :param follow_redirects: Whether to follow redirects. Defaults to "safe", which follows redirects but rejects those targeting internal/private IPs (SSRF protection).
+            Pass True to follow all redirects without restriction.
         :param max_redirects: Maximum number of redirects. Default 30, use -1 for unlimited.
         :param retries: Number of retry attempts. Defaults to 3.
         :param retry_delay: Number of seconds to wait between retry attempts. Defaults to 1 second.
@@ -371,7 +437,8 @@ async def bulk_get(
         :param headers: Headers to include in the request.
         :param cookies: Cookies to use in the request.
         :param timeout: Number of seconds to wait before timing out.
-        :param follow_redirects: Whether to follow redirects. Defaults to "safe", which follows redirects but rejects those targeting internal/private IPs (SSRF protection). Pass True to follow all redirects without restriction.
+        :param follow_redirects: Whether to follow redirects. Defaults to "safe", which follows redirects but rejects those targeting internal/private IPs (SSRF protection).
+            Pass True to follow all redirects without restriction.
         :param max_redirects: Maximum number of redirects. Default 30, use -1 for unlimited.
         :param retries: Number of retry attempts. Defaults to 3.
         :param retry_delay: Number of seconds to wait between retry attempts. Defaults to 1 second.
@@ -835,4 +902,6 @@ def serve(self, http: bool, host: str, port: int):
             description=self.bulk_stealthy_fetch.__doc__,
             structured_output=True,
         )
+        # Screenshot tool (returns image + url content blocks, not structured JSON)
+        server.add_tool(self.screenshot, title="screenshot", description=self.screenshot.__doc__)
         server.run(transport="stdio" if not http else "streamable-http")

tests/ai/test_ai_mcp.py

@@ -1,5 +1,12 @@
+import base64
+import struct
+from contextlib import contextmanager
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from threading import Thread
+
 import pytest
 import pytest_httpbin
+from mcp.types import ImageContent, TextContent
 
 from scrapling.core.ai import (
     ScraplingMCPServer,
@@ -197,6 +204,118 @@ async def test_open_session_duplicate_id_raises(self, server):
         await server.close_session("dupe")
 
 
+def _png_height(data: bytes) -> int:
+    """Read the height field from a PNG IHDR chunk."""
+    return struct.unpack(">I", data[20:24])[0]
+
+
+@contextmanager
+def _serve_html(body: bytes):
+    """Serve a fixed HTML body on localhost, yielding its URL."""
+
+    class _Handler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html; charset=utf-8")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+        def log_message(self, *args, **kwargs):
+            pass
+
+    server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
+    thread = Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield f"http://127.0.0.1:{server.server_address[1]}/"
+    finally:
+        server.shutdown()
+        server.server_close()
+
+
+@pytest_httpbin.use_class_based_httpbin
+class TestScreenshot:
+    """Test the screenshot tool"""
+
+    @pytest.fixture(scope="class")
+    def test_url(self, httpbin):
+        return f"{httpbin.url}/html"
+
+    @pytest.fixture
+    def server(self):
+        return ScraplingMCPServer()
+
+    @pytest.mark.asyncio
+    async def test_screenshot_png_with_dynamic_session(self, server, test_url):
+        """PNG screenshot via a dynamic session returns image and url content blocks"""
+        opened = await server.open_session(session_type="dynamic", headless=True)
+        try:
+            result = await server.screenshot(url=test_url, session_id=opened.session_id)
+            assert isinstance(result, list) and len(result) == 2
+            assert isinstance(result[0], ImageContent)
+            assert result[0].mimeType == "image/png"
+            assert isinstance(result[1], TextContent)
+            assert result[1].text == test_url
+        finally:
+            await server.close_session(opened.session_id)
+
+    @pytest.mark.asyncio
+    async def test_screenshot_jpeg_with_quality(self, server, test_url):
+        """JPEG screenshot with quality parameter via a dynamic session"""
+        opened = await server.open_session(session_type="dynamic", headless=True)
+        try:
+            result = await server.screenshot(url=test_url, session_id=opened.session_id, image_type="jpeg", quality=80)
+            assert isinstance(result[0], ImageContent)
+            assert result[0].mimeType == "image/jpeg"
+        finally:
+            await server.close_session(opened.session_id)
+
+    @pytest.mark.asyncio
+    async def test_screenshot_with_stealthy_session(self, server, test_url):
+        """PNG screenshot via a stealthy session"""
+        opened = await server.open_session(session_type="stealthy", headless=True)
+        try:
+            result = await server.screenshot(url=test_url, session_id=opened.session_id)
+            assert isinstance(result[0], ImageContent)
+            assert result[0].mimeType == "image/png"
+        finally:
+            await server.close_session(opened.session_id)
+
+    @pytest.mark.asyncio
+    async def test_screenshot_full_page_taller_than_viewport(self, server):
+        """full_page=True produces an image taller than the viewport-only capture"""
+        tall_html = b"<html><body><div style='height:5000px;background:#abc'></div></body></html>"
+        with _serve_html(tall_html) as tall_url:
+            opened = await server.open_session(session_type="dynamic", headless=True)
+            try:
+                viewport_result = await server.screenshot(url=tall_url, session_id=opened.session_id, full_page=False)
+                full_result = await server.screenshot(url=tall_url, session_id=opened.session_id, full_page=True)
+
+                viewport_png = base64.b64decode(viewport_result[0].data)
+                full_png = base64.b64decode(full_result[0].data)
+
+                assert _png_height(full_png) > _png_height(viewport_png)
+            finally:
+                await server.close_session(opened.session_id)
+
+    @pytest.mark.asyncio
+    async def test_screenshot_invalid_session_id_raises(self, server, test_url):
+        """Unknown session_id raises ValueError"""
+        with pytest.raises(ValueError, match="not found"):
+            await server.screenshot(url=test_url, session_id="does-not-exist")
+
+    @pytest.mark.asyncio
+    async def test_screenshot_quality_with_png_raises(self, server, test_url):
+        """quality is rejected when image_type is png"""
+        opened = await server.open_session(session_type="dynamic", headless=True)
+        try:
+            with pytest.raises(ValueError, match="quality"):
+                await server.screenshot(url=test_url, session_id=opened.session_id, image_type="png", quality=90)
+        finally:
+            await server.close_session(opened.session_id)
+
+
 class TestNormalizeCredentials:
     """Test the _normalize_credentials helper"""