fix(chrome-extension): keep recorder screenshots on their source tab

apps/chrome-extension/src/scripts/captureTabScreenshot.ts

@@ -0,0 +1,32 @@
+/// <reference types="chrome" />
+
+export async function captureTabScreenshot(
+  tab: chrome.tabs.Tab,
+): Promise<string | null> {
+  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);
+  }
+}

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');

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<typeof rs.fn>,
+  captureVisibleTab: ReturnType<typeof rs.fn>,
+) {
+  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);
+  });
+});