diff --git a/scrapling/core/ai.py b/scrapling/core/ai.py index 315733f..464d35b 100644 --- a/scrapling/core/ai.py +++ b/scrapling/core/ai.py @@ -3,7 +3,8 @@ from asyncio import gather 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 @@ from scrapling.core._types import ( ) 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 @@ class ScraplingMCPServer: 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 @@ class ScraplingMCPServer: :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 @@ class ScraplingMCPServer: :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 @@ class ScraplingMCPServer: 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") diff --git a/tests/ai/test_ai_mcp.py b/tests/ai/test_ai_mcp.py index 4806088..19e2f02 100644 --- a/tests/ai/test_ai_mcp.py +++ b/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 @@ class TestSessionManagement: 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"
" + 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"""