feat: add screenshot

meirk-brd · meirk-brd · commit 00541b233d47 · 2026-03-25T22:06:38.000+02:00
diff --git a/src/__tests__/browser/daemon.test.ts b/src/__tests__/browser/daemon.test.ts
@@ -45,6 +45,7 @@ class Mock_page extends EventEmitter {
     private current_title = 'Blank';
     private current_url = 'about:blank';
     private mock_context: Mock_context|null = null;
+    private screenshot_opts: {fullPage?: boolean}|null = null;
 
     constructor(mock_context?: Mock_context){
         super();
@@ -80,6 +81,13 @@ class Mock_page extends EventEmitter {
         return {status: ()=>200};
     }
 
+    async screenshot(opts?: {fullPage?: boolean}){
+        this.screenshot_opts = opts ?? {};
+        return Buffer.from(
+            opts?.fullPage ? 'full-page-image' : 'viewport-image'
+        );
+    }
+
     async evaluate(_fn: unknown, arg: {attr_name: string; selector?: string}){
         return {
             nodes: [
@@ -143,6 +151,10 @@ class Mock_page extends EventEmitter {
         this.closed = true;
         this.emit('close');
     }
+
+    last_screenshot_opts(){
+        return this.screenshot_opts;
+    }
 }
 
 class Mock_context extends EventEmitter {
@@ -536,6 +548,44 @@ describe('browser/daemon', ()=>{
         await daemon.stop();
     });
 
+    it('captures screenshots and returns the saved file path', async()=>{
+        const mock_browser = new Mock_browser();
+        const connect_over_cdp = vi.fn(async()=>mock_browser as unknown as Browser);
+        const daemon = new BrowserDaemon({
+            cdp_endpoint: 'wss://example.test',
+            daemon_dir: tmp_dir,
+            idle_timeout_ms: 0,
+            session_name: 'screenshot-actions',
+        }, {connect_over_cdp});
+
+        const output_path = path.join(tmp_dir, 'captures', 'page.png');
+        const screenshot = await daemon.handle_request({
+            id: 'screenshot-1',
+            action: 'screenshot',
+            params: {
+                base64: true,
+                full_page: true,
+                path: output_path,
+            },
+        });
+        expect(screenshot).toMatchObject({
+            id: 'screenshot-1',
+            success: true,
+            data: {
+                base64: Buffer.from('full-page-image').toString('base64'),
+                full_page: true,
+                mime_type: 'image/png',
+                path: output_path,
+            },
+        });
+
+        const page = mock_browser.first_context().first_page();
+        expect(page.last_screenshot_opts()).toEqual({fullPage: true});
+        expect(fs.readFileSync(output_path, 'utf8')).toBe('full-page-image');
+
+        await daemon.stop();
+    });
+
     it('starts and clears the keepalive interval with the daemon lifecycle', async()=>{
         vi.useFakeTimers();
 
diff --git a/src/__tests__/commands/browser.test.ts b/src/__tests__/commands/browser.test.ts
@@ -73,6 +73,7 @@ import {
     handle_browser_cookies,
     handle_browser_open,
     handle_browser_reload,
+    handle_browser_screenshot,
     handle_browser_sessions,
     handle_browser_snapshot,
     handle_browser_status,
@@ -216,6 +217,39 @@ describe('commands/browser', ()=>{
         );
     });
 
+    it('captures a screenshot for an active browser session', async()=>{
+        mocks.send_command.mockResolvedValue({
+            success: true,
+            data: {
+                full_page: true,
+                mime_type: 'image/png',
+                path: '/tmp/browser-shot.png',
+            },
+        });
+
+        await handle_browser_screenshot('/tmp/browser-shot.png', {
+            fullPage: true,
+            session: 'shop',
+        });
+
+        expect(mocks.send_command).toHaveBeenCalledWith(
+            'shop',
+            expect.objectContaining({
+                action: 'screenshot',
+                params: {
+                    base64: false,
+                    full_page: true,
+                    path: '/tmp/browser-shot.png',
+                },
+            }),
+            {daemon_dir: undefined, timeout_ms: undefined}
+        );
+        expect(mocks.print).toHaveBeenCalledWith(
+            '/tmp/browser-shot.png',
+            {output: undefined}
+        );
+    });
+
     it('rejects invalid snapshot depth before sending the command', async()=>{
         await expect(handle_browser_snapshot({depth: 'abc', session: 'shop'})).rejects.toThrow(
             'fail:Snapshot depth must be a non-negative integer.'
@@ -454,6 +488,52 @@ describe('commands/browser', ()=>{
         );
     });
 
+    it('parses screenshot flags and forwards the screenshot params', async()=>{
+        mocks.send_command.mockResolvedValue({
+            success: true,
+            data: {
+                base64: 'aW1hZ2U=',
+                full_page: true,
+                mime_type: 'image/png',
+                path: '/tmp/browser-shot.png',
+            },
+        });
+        const command = create_browser_command();
+        command.exitOverride();
+
+        await command.parseAsync([
+            'screenshot',
+            '/tmp/browser-shot.png',
+            '--session',
+            'shop',
+            '--full-page',
+            '--base64',
+            '--json',
+        ], {from: 'user'});
+
+        expect(mocks.send_command).toHaveBeenCalledWith(
+            'shop',
+            expect.objectContaining({
+                action: 'screenshot',
+                params: {
+                    base64: true,
+                    full_page: true,
+                    path: '/tmp/browser-shot.png',
+                },
+            }),
+            {daemon_dir: undefined, timeout_ms: undefined}
+        );
+        expect(mocks.print).toHaveBeenCalledWith(
+            {
+                base64: 'aW1hZ2U=',
+                full_page: true,
+                mime_type: 'image/png',
+                path: '/tmp/browser-shot.png',
+            },
+            {json: true, output: undefined, pretty: undefined}
+        );
+    });
+
     it('rejects open-only flags on status through browser-group parsing', async()=>{
         const command = create_browser_command();
         command.exitOverride();
@@ -498,6 +578,20 @@ describe('commands/browser', ()=>{
         expect(mocks.send_command).not.toHaveBeenCalled();
     });
 
+    it('rejects screenshot-only flags outside the screenshot command', async()=>{
+        const command = create_browser_command();
+        command.exitOverride();
+
+        await expect(command.parseAsync([
+            'status',
+            '--full-page',
+        ], {from: 'user'})).rejects.toThrow(
+            'fail:--full-page is not supported by "brightdata browser status".'
+        );
+
+        expect(mocks.send_command).not.toHaveBeenCalled();
+    });
+
     it('rejects not-yet-implemented global flags on open', async()=>{
         const command = create_browser_command();
         command.exitOverride();
diff --git a/src/browser/daemon.ts b/src/browser/daemon.ts
@@ -5,6 +5,7 @@ import path from 'path';
 import {chromium} from 'playwright-core';
 import {clear_connection_state, ensure_connected as ensure_browser_connected} from './connection';
 import {parse_daemon_request} from './ipc';
+import {take_screenshot} from './screenshot';
 import {capture_snapshot} from './snapshot';
 import type {
     Browser,
@@ -505,6 +506,8 @@ class BrowserDaemon {
             return {alive: true, connected: this.state.connected};
         case 'reload':
             return this.handle_reload();
+        case 'screenshot':
+            return this.handle_screenshot(request.params);
         case 'snapshot':
             return this.handle_snapshot(request.params);
         case 'status':
@@ -620,6 +623,39 @@ class BrowserDaemon {
         };
     }
 
+    private async handle_screenshot(params: Json_object|undefined){
+        const base64 = params?.['base64'];
+        const full_page = params?.['full_page'];
+        const file_path = params?.['path'];
+
+        if (base64 !== undefined && typeof base64 != 'boolean')
+        {
+            throw new Error(
+                'Screenshot "base64" parameter must be a boolean when provided.'
+            );
+        }
+        if (full_page !== undefined && typeof full_page != 'boolean')
+        {
+            throw new Error(
+                'Screenshot "full_page" parameter must be a boolean when provided.'
+            );
+        }
+        if (file_path !== undefined
+            && (typeof file_path != 'string' || !file_path.trim()))
+        {
+            throw new Error(
+                'Screenshot "path" parameter must be a non-empty string when provided.'
+            );
+        }
+
+        const page = await this.ensure_connected();
+        return await take_screenshot(page, {
+            base64: base64 === true,
+            full_page: full_page === true,
+            path: typeof file_path == 'string' ? file_path.trim() : undefined,
+        });
+    }
+
     private handle_network(){
         return {
             requests: Array.from(this.state.requests.entries()).map(
diff --git a/src/browser/screenshot.ts b/src/browser/screenshot.ts
@@ -0,0 +1,71 @@
+import crypto from 'crypto';
+import fs from 'fs';
+import os from 'os';
+import path from 'path';
+import type {Page} from 'playwright-core';
+
+const SCREENSHOT_MIME_TYPE = 'image/png';
+const SCREENSHOT_TMP_DIR = 'brightdata-cli';
+const SCREENSHOT_TMP_PREFIX = 'browser-screenshot-';
+
+type Screenshot_capture_opts = {
+    base64?: boolean;
+    full_page?: boolean;
+    path?: string;
+};
+
+type Screenshot_capture_result = {
+    base64?: string;
+    full_page: boolean;
+    mime_type: string;
+    path: string;
+};
+
+const normalize_screenshot_path = (file_path: string|undefined): string|undefined=>{
+    if (file_path === undefined)
+        return undefined;
+    const normalized = file_path.trim();
+    if (!normalized)
+        throw new Error('Screenshot path cannot be empty.');
+    return path.resolve(normalized);
+};
+
+const create_temp_screenshot_path = (): string=>{
+    const token = crypto.randomBytes(6).toString('hex');
+    return path.join(
+        os.tmpdir(),
+        SCREENSHOT_TMP_DIR,
+        `${SCREENSHOT_TMP_PREFIX}${Date.now()}-${token}.png`
+    );
+};
+
+const take_screenshot = async(
+    page: Page,
+    opts: Screenshot_capture_opts = {},
+): Promise<Screenshot_capture_result>=>{
+    const full_page = opts.full_page === true;
+    const output_path = normalize_screenshot_path(opts.path)
+        ?? create_temp_screenshot_path();
+    const buffer = await page.screenshot({fullPage: full_page});
+
+    fs.mkdirSync(path.dirname(output_path), {recursive: true});
+    fs.writeFileSync(output_path, buffer);
+
+    return {
+        base64: opts.base64 === true ? buffer.toString('base64') : undefined,
+        full_page,
+        mime_type: SCREENSHOT_MIME_TYPE,
+        path: output_path,
+    };
+};
+
+export {
+    create_temp_screenshot_path,
+    normalize_screenshot_path,
+    SCREENSHOT_MIME_TYPE,
+    take_screenshot,
+};
+export type {
+    Screenshot_capture_opts,
+    Screenshot_capture_result,
+};
diff --git a/src/commands/browser.ts b/src/commands/browser.ts
