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);
+ });
+});