NTP omnibar: keep Add Screenshot enabled while a capture is pending

Native now answers a pending capture as cancelled when a new one replaces it,
so the page no longer greys the screenshot rows while a request is pending.
Only the latest request may show the capture error.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

special-pages/pages/new-tab/app/omnibar/components/Omnibar.js

@@ -502,11 +502,7 @@ function AiChatContent({
                                             ? {
                                                   modes: screenshotModes,
                                                   onCapture: screenshotCapture.capture,
-                                                  disabled:
-                                                      blocksPrompt ||
-                                                      !canAttachImages ||
-                                                      screenshotCapture.capturing ||
-                                                      imageState.imageUploadDisabled,
+                                                  disabled: blocksPrompt || !canAttachImages || imageState.imageUploadDisabled,
                                               }
                                             : null
                                     }

special-pages/pages/new-tab/app/omnibar/components/chat-tools/image-attachment/useScreenshotCapture.js

@@ -46,59 +46,64 @@ export function useScreenshotCapture({ imageState, canAttachImages, processOther
     const ntp = useMessaging();
     const { captureScreenshot } = useContext(OmnibarContext);
     const [captureError, setCaptureError] = useState(false);
-    const [capturing, setCapturing] = useState(false);
 
     // A capture can stay pending for as long as the user is in the native picker; read the image
     // state as it is when the reply arrives, not as it was when the capture started.
     const imageStateRef = useRef(imageState);
     imageStateRef.current = imageState;
 
+    // Native answers a request as cancelled when a newer one replaces it, so replies can arrive
+    // after a later capture has started. Only the latest request may show an error.
+    const latestRequestRef = useRef(0);
+
     const takenNames = () => new Set(imageStateRef.current.attachedImages.map((img) => img.fileName));
 
-    /** @param {ScreenshotMode} mode */
+    /**
+     * Starts a capture. The rows stay enabled while it is pending: choosing again sends a new
+     * request, which native lets replace the pending one.
+     * @param {ScreenshotMode} mode
+     */
     const capture = async (mode) => {
-        if (capturing) return;
+        const requestId = ++latestRequestRef.current;
+        const isLatest = () => requestId === latestRequestRef.current;
         setCaptureError(false);
-        setCapturing(true);
+
+        /** @type {import('../../../../../types/new-tab.js').CaptureScreenshotResponse} */
+        let response;
+        try {
+            response = await captureScreenshot(mode);
+        } catch (err) {
+            console.warn('omnibar_captureScreenshot failed', err);
+            if (isLatest()) setCaptureError(true);
+            return;
+        }
+        // Native reports its own capture failures; the page only shows the message.
+        if (response.error) {
+            if (isLatest()) setCaptureError(true);
+            return;
+        }
+        if (!response.image) return; // cancelled, or replaced by a newer request
+
+        // An image is the user's capture whichever request it answers, so it is always attached.
+        const { image } = response;
+        /** @type {File} */
+        let file;
         try {
-            /** @type {import('../../../../../types/new-tab.js').CaptureScreenshotResponse} */
-            let response;
-            try {
-                response = await captureScreenshot(mode);
-            } catch (err) {
-                console.warn('omnibar_captureScreenshot failed', err);
-                setCaptureError(true);
-                return;
-            }
-            // Native reports its own capture failures; the page only shows the message.
-            if (response.error) {
-                setCaptureError(true);
-                return;
-            }
-            if (!response.image) return; // cancelled
-
-            const { image } = response;
-            /** @type {File} */
-            let file;
-            try {
-                file = screenshotToFile(image, uniqueFileName(t('omnibar_screenshotFileName'), extensionFor(image.format), takenNames()));
-            } catch (err) {
-                console.warn('Screenshot rejected: invalid image data');
-                ntp.telemetryEvent({ attributes: { name: 'omnibar_screenshot_failed', value: { reason: 'failed' } } });
-                setCaptureError(true);
-                return;
-            }
-            const result = await imageStateRef.current.processFiles([file], {
-                maxDimension: SCREENSHOT_MAX_DIMENSION,
-                source: 'screenshot',
-            });
-            if (result.added > 0) {
-                ntp.telemetryEvent({ attributes: { name: 'omnibar_screenshot_taken', value: { kind: image.kind } } });
-            } else if (result.rejected > 0) {
-                ntp.telemetryEvent({ attributes: { name: 'omnibar_screenshot_failed', value: { reason: 'failed' } } });
-            }
-        } finally {
-            setCapturing(false);
+            file = screenshotToFile(image, uniqueFileName(t('omnibar_screenshotFileName'), extensionFor(image.format), takenNames()));
+        } catch (err) {
+            console.warn('Screenshot rejected: invalid image data');
+            ntp.telemetryEvent({ attributes: { name: 'omnibar_screenshot_failed', value: { reason: 'failed' } } });
+            if (isLatest()) setCaptureError(true);
+            return;
+        }
+        const result = await imageStateRef.current.processFiles([file], {
+            maxDimension: SCREENSHOT_MAX_DIMENSION,
+            source: 'screenshot',
+        });
+        if (result.added > 0) {
+            ntp.telemetryEvent({ attributes: { name: 'omnibar_screenshot_taken', value: { kind: image.kind } } });
+        } else if (result.rejected > 0) {
+            ntp.telemetryEvent({ attributes: { name: 'omnibar_screenshot_failed', value: { reason: 'failed' } } });
         }
     };
 
@@ -160,7 +165,6 @@ export function useScreenshotCapture({ imageState, canAttachImages, processOther
 
     return {
         capture,
-        capturing,
         captureError,
         clearCaptureError: () => setCaptureError(false),
         handlePaste,

special-pages/pages/new-tab/app/omnibar/integration-tests/omnibar-screenshot.spec.js

@@ -141,15 +141,14 @@ test.describe('omnibar screenshot menu', () => {
         await expect(omnibar.addScreenshotMenuItem()).toHaveAttribute('aria-disabled', 'true');
     });
 
-    test('greys out "Add Screenshot" while a capture is pending', async ({ page }, workerInfo) => {
+    test('keeps "Add Screenshot" enabled while a capture is pending', async ({ page }, workerInfo) => {
         const { omnibar } = await setup(page, workerInfo, { 'omnibar.screenshotModes': BOTH_MODES, 'omnibar.screenshotDelay': '1500' });
 
         await omnibar.captureScreenshot('Drag to Select');
         await omnibar.attachMenuButton().click();
-        await expect(omnibar.addScreenshotMenuItem()).toHaveAttribute('aria-disabled', 'true');
-
-        await expect(omnibar.imagePreviews()).toHaveCount(1);
-        await expect(omnibar.addScreenshotMenuItem()).not.toHaveAttribute('aria-disabled', 'true');
+        await expect(omnibar.addScreenshotMenuItem()).not.toHaveAttribute('aria-disabled');
+        await omnibar.addScreenshotMenuItem().click();
+        await expect(omnibar.screenshotModeItem('Drag to Select')).not.toHaveAttribute('aria-disabled');
     });
 
     test('opens and closes from the keyboard', async ({ page }, workerInfo) => {
@@ -226,6 +225,25 @@ test.describe('omnibar screenshot capture', () => {
             .toEqual([{ name: 'omnibar_image_removed', value: { source: 'screenshot' } }, { name: 'omnibar_screenshot_removed' }]);
     });
 
+    test('a new Drag to Select replaces a pending one, whose cancelled reply shows nothing', async ({ page }, workerInfo) => {
+        const { omnibar } = await setup(page, workerInfo, { 'omnibar.screenshotModes': BOTH_MODES, 'omnibar.screenshotDelay': '1000' });
+
+        await omnibar.captureScreenshot('Drag to Select');
+        await omnibar.expectMethodCallCount('omnibar_captureScreenshot', 1);
+        await omnibar.captureScreenshot('Drag to Select');
+        await omnibar.expectMethodCallCount('omnibar_captureScreenshot', 2);
+
+        // The first request is answered as cancelled straight away; only the second one attaches.
+        await expect(omnibar.imagePreviews()).toHaveCount(1);
+        await expect(page.getByRole('alert')).toHaveCount(0);
+        await expect
+            .poll(() => omnibar.telemetryEvents())
+            .toEqual([
+                { name: 'omnibar_image_attached', value: { source: 'screenshot' } },
+                { name: 'omnibar_screenshot_taken', value: { kind: 'selection' } },
+            ]);
+    });
+
     test('numbers repeated screenshots so each one attaches', async ({ page }, workerInfo) => {
         const { omnibar } = await setup(page, workerInfo, { 'omnibar.screenshotModes': BOTH_MODES });
 

special-pages/pages/new-tab/app/omnibar/mocks/omnibar.mock-transport.js

@@ -259,6 +259,12 @@ export function omnibarMockTransport() {
 
     /** @type {Set<string>} Tracks deleted chats so re-fetches exclude them */
     const deletedChatIds = new Set();
+    /**
+     * Settles the capture request that is still waiting, if any. A new request answers it as
+     * cancelled, as native does for Drag to Select (the mock does it for either mode).
+     * @type {((reply: import('../../../types/new-tab.js').CaptureScreenshotResponse) => void) | null}
+     */
+    let settlePendingCapture = null;
 
     return new TestTransportConfig({
         notify(_msg) {
@@ -601,9 +607,20 @@ export function omnibarMockTransport() {
                     return getMockOpenTabs(openTabsCount >= 0 ? openTabsCount : undefined);
                 }
                 case 'omnibar_captureScreenshot': {
+                    settlePendingCapture?.({});
                     // Simulates the user spending time in the native picker or selection overlay.
                     const delay = parseInt(url.searchParams.get('omnibar.screenshotDelay') ?? '', 10);
-                    await new Promise((resolve) => setTimeout(resolve, delay >= 0 ? delay : window.__playwright_01 ? 0 : 600));
+                    /** @type {import('../../../types/new-tab.js').CaptureScreenshotResponse | null} */
+                    const replaced = await new Promise((resolve) => {
+                        /** @param {import('../../../types/new-tab.js').CaptureScreenshotResponse | null} reply */
+                        const settle = (reply) => {
+                            if (settlePendingCapture === settle) settlePendingCapture = null;
+                            resolve(reply);
+                        };
+                        settlePendingCapture = settle;
+                        setTimeout(() => settle(null), delay >= 0 ? delay : window.__playwright_01 ? 0 : 600);
+                    });
+                    if (replaced) return replaced;
                     const override = window.__playwright_01?.mockResponses?.omnibar_captureScreenshot;
                     if (override) return /** @type {import('../../../types/new-tab.js').CaptureScreenshotResponse} */ (override);
                     // `omnibar.screenshotResult=error|cancel|invalid`; anything else returns a capture

special-pages/pages/new-tab/app/omnibar/omnibar.md

@@ -38,7 +38,7 @@ title: Omnibar Widget
   - `enableVoiceChatAccess` — when true and the input is empty, replaces the AI chat submit button with a 1-click voice-chat button. Click/Enter sends `omnibar_submitChat` with an empty `chat` and `mode: "voice-mode"` — native handles the voice handoff (default `false`)
   - `enableAskAiSuggestion` — when `false`, hides the inline "Ask Duck.ai: <query>" entry in the suggestions dropdown. Missing/undefined is treated as `true` (default `true`). Does not affect the Duck.ai mode pill or any other AI affordance — those remain governed by `enableAi`
   - `enableAttachTabs` — when `true`, the omnibar shows the page context entry point and accepts `@` mentions for attaching open tabs as context. Requires native to handle `omnibar_getOpenTabs` and `omnibar_getTabContent` (default `false`).
-  - `screenshotModes` — capture modes (`"dragToSelect"`, `"selectWindowOrDisplay"`) listed, in order, under "Add Screenshot" in the paperclip menu. Absent or empty hides the screenshot UI. Requires native to handle `omnibar_captureScreenshot`. The row is disabled while a capture is pending, at the image cap, when the model cannot take images, or when the prompt is blocked.
+  - `screenshotModes` — capture modes (`"dragToSelect"`, `"selectWindowOrDisplay"`) listed, in order, under "Add Screenshot" in the paperclip menu. Absent or empty hides the screenshot UI. Requires native to handle `omnibar_captureScreenshot`. The row is disabled at the image cap, when the model cannot take images, or when the prompt is blocked; it stays enabled while a capture is pending.
   - `enablePastedAttachments` — when `true`, pasting into the Duck.ai prompt attaches copied images and files (see [Paste](#paste)). When `false` or absent, paste is left to the browser (text only) (default `false`).
   - `aiModelSections` — array of model sections for the model selector. Each model may include `supportedReasoningEffort` (e.g. `["none", "low", "medium"]`) to surface the reasoning picker
   - `selectedModelId` — the user's persisted model choice
@@ -61,7 +61,7 @@ title: Omnibar Widget
 
 ### `omnibar_captureScreenshot`
 - {@link "NewTab Messages".OmnibarCaptureScreenshotRequest}
-- Sent when the user picks a mode from the "Add Screenshot" submenu. Native runs the capture (selection overlay or window/display picker) and replies when the user finishes or cancels, so the request can stay pending for as long as that takes. The page sends one request at a time.
+- Sent when the user picks a mode from the "Add Screenshot" submenu. Native runs the capture (selection overlay or window/display picker) and replies when the user finishes or cancels, so the request can stay pending for as long as that takes. The user can choose a mode again while one is pending: native answers the older request as cancelled and serves the new one (a Drag to Select replaces a pending one; a Select Window or Display request while the picker is open is answered as cancelled). Only the latest request's `error` is shown.
 - requires `mode`, one of the configured `screenshotModes`.
 - returns {@link "NewTab Messages".CaptureScreenshotResponse}:
   - `image` — base64 `data` (no data-URL prefix), `format` (`png` or `jpeg`) and `kind` (`selection`, `screen` or `window`). Native scales it to at most 1024px on the long side; the page keeps that size and adds it as an image chip named "Screenshot" (numbered on repeats).

special-pages/pages/new-tab/messages/omnibar_captureScreenshot.request.json

@@ -1,7 +1,7 @@
 {
   "$schema": "http://json-schema.org/draft-07/schema#",
   "title": "Capture Screenshot Params",
-  "description": "Asks native to capture a screenshot for the Duck.ai prompt. The request stays pending while the user is in the native picker or selection overlay; native allows one capture at a time.",
+  "description": "Asks native to capture a screenshot for the Duck.ai prompt. The request stays pending while the user is in the native picker or selection overlay. A new request while one is pending makes native answer the older one as cancelled.",
   "type": "object",
   "required": ["mode"],
   "properties": {

special-pages/pages/new-tab/types/new-tab.ts

@@ -1686,7 +1686,7 @@ export interface OmnibarCaptureScreenshotRequest {
   result: CaptureScreenshotResponse;
 }
 /**
- * Asks native to capture a screenshot for the Duck.ai prompt. The request stays pending while the user is in the native picker or selection overlay; native allows one capture at a time.
+ * Asks native to capture a screenshot for the Duck.ai prompt. The request stays pending while the user is in the native picker or selection overlay. A new request while one is pending makes native answer the older one as cancelled.
  */
 export interface CaptureScreenshotParams {
   mode: ScreenshotMode;