feat: Add enhanced screenshot API

Author: avocardio

browser_use/browser/session.py

@@ -2255,3 +2255,129 @@ async def cdp_client_for_node(self, node: EnhancedDOMTreeNode) -> CDPSession:
 				self.logger.debug(f'Failed to get CDP client for target {node.target_id}: {e}, using main session')
 
 		return await self.get_or_create_cdp_session()
+
+	async def take_screenshot(self, path: str | None = None, full_page: bool = False, format: str = 'png', quality: int | None = None, clip: dict | None = None) -> bytes:
+		"""Take a screenshot using CDP.
+		
+		Args:
+			path: Optional file path to save screenshot
+			full_page: Capture entire scrollable page beyond viewport
+			format: Image format ('png', 'jpeg', 'webp')
+			quality: Quality 0-100 for JPEG format
+			clip: Region to capture {'x': int, 'y': int, 'width': int, 'height': int}
+		
+		Returns:
+			Screenshot data as bytes
+		"""
+		import base64
+		from cdp_use.cdp.page import CaptureScreenshotParameters
+
+		cdp_session = await self.get_or_create_cdp_session()
+
+		# Build parameters dict explicitly to satisfy TypedDict expectations
+		params: CaptureScreenshotParameters = {
+			'format': format,
+			'captureBeyondViewport': full_page,
+		}
+
+		if quality is not None and format == 'jpeg':
+			params['quality'] = quality
+
+		if clip:
+			params['clip'] = {
+				'x': clip['x'],
+				'y': clip['y'],
+				'width': clip['width'],
+				'height': clip['height'],
+				'scale': 1,
+			}
+
+		params = CaptureScreenshotParameters(**params)
+
+		result = await cdp_session.cdp_client.send.Page.captureScreenshot(
+			params=params,
+			session_id=cdp_session.session_id
+		)
+
+		if not result or 'data' not in result:
+			raise Exception('Screenshot failed - no data returned')
+
+		screenshot_data = base64.b64decode(result['data'])
+
+		if path:
+			with open(path, 'wb') as f:
+				f.write(screenshot_data)
+
+		return screenshot_data
+
+	async def screenshot_element(
+		self,
+		selector: str,
+		path: str | None = None,
+		format: str = 'png',
+		quality: int | None = None,
+	) -> bytes:
+		"""Take a screenshot of a specific element.
+
+		Args:
+			selector: CSS selector for the element
+			path: Optional file path to save screenshot
+			format: Image format ('png', 'jpeg', 'webp')
+			quality: Quality 0-100 for JPEG format
+
+		Returns:
+			Screenshot data as bytes
+		"""
+
+		bounds = await self._get_element_bounds(selector)
+		if not bounds:
+			raise ValueError(f"Element '{selector}' not found or has no bounds")
+
+		return await self.take_screenshot(
+			path=path,
+			format=format,
+			quality=quality,
+			clip=bounds,
+		)
+
+	async def _get_element_bounds(self, selector: str) -> dict | None:
+		"""Get element bounding box using CDP."""
+
+		cdp_session = await self.get_or_create_cdp_session()
+
+		# Get document
+		doc = await cdp_session.cdp_client.send.DOM.getDocument(
+			params={'depth': 1},
+			session_id=cdp_session.session_id
+		)
+
+		# Query selector
+		node_result = await cdp_session.cdp_client.send.DOM.querySelector(
+			params={
+				'nodeId': doc['root']['nodeId'],
+				'selector': selector
+			},
+			session_id=cdp_session.session_id
+		)
+
+		node_id = node_result.get('nodeId')
+		if not node_id:
+			return None
+
+		# Get bounding box
+		box_result = await cdp_session.cdp_client.send.DOM.getBoxModel(
+			params={'nodeId': node_id},
+			session_id=cdp_session.session_id
+		)
+
+		box_model = box_result.get('model')
+		if not box_model:
+			return None
+
+		content = box_model['content']
+		return {
+			'x': min(content[0], content[2], content[4], content[6]),
+			'y': min(content[1], content[3], content[5], content[7]),
+			'width': max(content[0], content[2], content[4], content[6]) - min(content[0], content[2], content[4], content[6]),
+			'height': max(content[1], content[3], content[5], content[7]) - min(content[1], content[3], content[5], content[7])
+		}

tests/ci/test_browser_session_take_screenshot.py

@@ -0,0 +1,33 @@
+import pytest
+
+from browser_use.browser.events import NavigateToUrlEvent
+from browser_use.browser.profile import BrowserProfile
+from browser_use.browser.session import BrowserSession
+
+
+@pytest.fixture(scope="function")
+async def browser_session():
+    session = BrowserSession(browser_profile=BrowserProfile(headless=True))
+    await session.start()
+    yield session
+    await session.kill()
+
+
+@pytest.mark.asyncio
+async def test_basic_screenshots(browser_session: BrowserSession, httpserver):
+    """Navigate to a local page and ensure screenshot helpers return bytes."""
+
+    html = """
+    <html><body><h1 id='title'>Hello</h1><p>Screenshot demo.</p></body></html>
+    """
+    httpserver.expect_request("/demo").respond_with_data(html, content_type="text/html")
+    url = httpserver.url_for("/demo")
+
+    nav = browser_session.event_bus.dispatch(NavigateToUrlEvent(url=url, new_tab=False))
+    await nav
+
+    data = await browser_session.take_screenshot(full_page=False)
+    assert data, "Viewport screenshot returned no data"
+
+    element = await browser_session.screenshot_element("h1")
+    assert element, "Element screenshot returned no data"