diff --git a/apps/chrome-extension/src/scripts/captureTabScreenshot.ts b/apps/chrome-extension/src/scripts/captureTabScreenshot.ts new file mode 100644 index 0000000000..e5ed2b312a --- /dev/null +++ b/apps/chrome-extension/src/scripts/captureTabScreenshot.ts @@ -0,0 +1,32 @@ +/// + +export async function captureTabScreenshot( + tab: chrome.tabs.Tab, +): Promise { + if (tab.id === undefined || tab.windowId === undefined) return null; + + const isSourceTabActive = async () => + (await chrome.tabs.query({ active: true, windowId: tab.windowId }))[0] + ?.id === tab.id; + + let sourceTabWasDeactivated = false; + const onActivated = (activeInfo: chrome.tabs.TabActiveInfo) => { + if (activeInfo.windowId === tab.windowId && activeInfo.tabId !== tab.id) { + sourceTabWasDeactivated = true; + } + }; + chrome.tabs.onActivated.addListener(onActivated); + + try { + if (!(await isSourceTabActive())) return null; + + const screenshot = await chrome.tabs.captureVisibleTab(tab.windowId, { + format: 'png', + }); + + const stillActive = await isSourceTabActive(); + return !sourceTabWasDeactivated && stillActive ? screenshot : null; + } finally { + chrome.tabs.onActivated.removeListener(onActivated); + } +} diff --git a/apps/chrome-extension/src/scripts/worker.ts b/apps/chrome-extension/src/scripts/worker.ts index 98da963c9d..78a42f9b21 100644 --- a/apps/chrome-extension/src/scripts/worker.ts +++ b/apps/chrome-extension/src/scripts/worker.ts @@ -5,6 +5,7 @@ import { uuid } from '@midscene/shared/utils'; import { BridgeConnector, type BridgeStatus } from '../utils/bridgeConnector'; import { registerAlarmListener, safeSetupKeepalive } from '../utils/keepalive'; import { workerMessageTypes } from '../utils/workerMessageTypes'; +import { captureTabScreenshot } from './captureTabScreenshot'; // save screenshot interface WorkerRequestSaveContext { @@ -438,21 +439,12 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => { // Handle screenshot capture request if (request.action === 'captureScreenshot') { if (sender.tab && sender.tab.id !== undefined) { - chrome.tabs.captureVisibleTab( - sender.tab.windowId, - { format: 'png' }, - (dataUrl) => { - if (chrome.runtime.lastError) { - console.error( - '[ServiceWorker] Failed to capture screenshot:', - chrome.runtime.lastError, - ); - sendResponse(null); - } else { - sendResponse(dataUrl); - } - }, - ); + captureTabScreenshot(sender.tab) + .then(sendResponse) + .catch((error) => { + console.error('[ServiceWorker] Failed to capture screenshot:', error); + sendResponse(null); + }); return true; // Keep the message channel open for async response } else { console.error('[ServiceWorker] No valid tab for screenshot capture'); diff --git a/apps/chrome-extension/tests/capture-tab-screenshot.test.ts b/apps/chrome-extension/tests/capture-tab-screenshot.test.ts new file mode 100644 index 0000000000..bd360d4610 --- /dev/null +++ b/apps/chrome-extension/tests/capture-tab-screenshot.test.ts @@ -0,0 +1,85 @@ +import { afterEach, describe, expect, it, rs } from '@rstest/core'; +import { captureTabScreenshot } from '../src/scripts/captureTabScreenshot'; + +const sourceTab = { id: 1, windowId: 10 } as chrome.tabs.Tab; + +function stubChrome( + query: ReturnType, + captureVisibleTab: ReturnType, +) { + const addListener = rs.fn(); + const removeListener = rs.fn(); + rs.stubGlobal('chrome', { + tabs: { + query, + captureVisibleTab, + onActivated: { addListener, removeListener }, + }, + }); + return { addListener, removeListener }; +} + +describe('captureTabScreenshot', () => { + afterEach(() => { + rs.unstubAllGlobals(); + }); + + it('captures the requesting tab while it remains active', async () => { + const query = rs.fn().mockResolvedValue([{ id: 1 }]); + const captureVisibleTab = rs + .fn() + .mockResolvedValue('data:image/png;base64,a'); + const { addListener, removeListener } = stubChrome( + query, + captureVisibleTab, + ); + + await expect(captureTabScreenshot(sourceTab)).resolves.toBe( + 'data:image/png;base64,a', + ); + expect(query).toHaveBeenCalledTimes(2); + expect(query).toHaveBeenCalledWith({ active: true, windowId: 10 }); + expect(captureVisibleTab).toHaveBeenCalledWith(10, { format: 'png' }); + expect(removeListener).toHaveBeenCalledWith(addListener.mock.calls[0][0]); + }); + + it('does not capture another tab when the source is already inactive', async () => { + const query = rs.fn().mockResolvedValue([{ id: 2 }]); + const captureVisibleTab = rs.fn(); + stubChrome(query, captureVisibleTab); + + await expect(captureTabScreenshot(sourceTab)).resolves.toBeNull(); + expect(captureVisibleTab).not.toHaveBeenCalled(); + }); + + it('discards a capture if the active tab changes during capture', async () => { + const query = rs + .fn() + .mockResolvedValueOnce([{ id: 1 }]) + .mockResolvedValueOnce([{ id: 2 }]); + const captureVisibleTab = rs.fn().mockResolvedValue('other tab screenshot'); + stubChrome(query, captureVisibleTab); + + await expect(captureTabScreenshot(sourceTab)).resolves.toBeNull(); + expect(captureVisibleTab).toHaveBeenCalledTimes(1); + }); + + it('discards a capture even if the source tab becomes active again', async () => { + const query = rs.fn().mockResolvedValue([{ id: 1 }]); + let onActivated: + | ((activeInfo: chrome.tabs.TabActiveInfo) => void) + | undefined; + const captureVisibleTab = rs.fn().mockImplementation(async () => { + onActivated?.({ tabId: 2, windowId: 10 }); + onActivated?.({ tabId: 1, windowId: 10 }); + return 'other tab screenshot'; + }); + const { addListener } = stubChrome(query, captureVisibleTab); + addListener.mockImplementation((listener) => { + onActivated = listener; + }); + + await expect(captureTabScreenshot(sourceTab)).resolves.toBeNull(); + expect(query).toHaveBeenCalledTimes(2); + }); +});