fix(desktop): preview markup survives page reloads

apps/desktop/src/preview/GuestProtocol.ts

@@ -2,6 +2,7 @@ export const START_PICK_CHANNEL = "preview:start-pick";
 export const CANCEL_PICK_CHANNEL = "preview:cancel-pick";
 export const ELEMENT_PICKED_CHANNEL = "preview:element-picked";
 export const ANNOTATION_CAPTURED_CHANNEL = "preview:annotation-captured";
+export const ANNOTATION_DRAFT_CHANNEL = "preview:annotation-draft";
 export const ANNOTATION_THEME_CHANNEL = "preview:annotation-theme";
 export const HUMAN_INPUT_CHANNEL = "preview:human-input";
 export const MOUSE_NAVIGATE_CHANNEL = "preview:mouse-navigate";

apps/desktop/src/preview/Manager.test.ts

@@ -3857,6 +3857,98 @@ describe("PreviewManager", () => {
     ),
   );
 
+  effectIt.effect("keeps the markup when the page reloads during element picking", () =>
+    withManager((manager) =>
+      Effect.gen(function* () {
+        const listeners = new Map<string, (...args: unknown[]) => void>();
+        let onDraft: ((event: unknown, value: unknown) => void) | undefined;
+        let currentUrl = "https://example.com/app#section";
+        fromId.mockReturnValue({
+          id: 42,
+          isDestroyed: () => false,
+          getType: () => "webview",
+          getURL: () => currentUrl,
+          getTitle: () => "Example",
+          isLoading: () => false,
+          isFocused: () => true,
+          getZoomFactor: () => 1,
+          setZoomFactor: vi.fn(),
+          setAudioMuted: vi.fn(),
+          isCurrentlyAudible: () => false,
+          on: vi.fn((event: string, listener: (...args: unknown[]) => void) => {
+            listeners.set(event, listener);
+          }),
+          once: vi.fn(),
+          off: vi.fn(),
+          ipc: {
+            on: vi.fn((channel: string, listener: typeof onDraft) => {
+              if (channel === "preview:annotation-draft") onDraft = listener;
+            }),
+            off: vi.fn(),
+            removeListener: vi.fn(),
+          },
+          send: webviewSend,
+          navigationHistory: { canGoBack: () => false, canGoForward: () => false },
+          setIgnoreMenuShortcuts: vi.fn(),
+          setWindowOpenHandler: vi.fn(),
+          debugger: {
+            isAttached: () => false,
+            attach: vi.fn(),
+            sendCommand: vi.fn(async () => undefined),
+            on: vi.fn(),
+            off: vi.fn(),
+          },
+        } as never);
+        const draft = {
+          comment: "Tighten this spacing",
+          tool: "marquee",
+          elements: [
+            {
+              id: "element_1",
+              selector: ":root > body:nth-of-type(1) > main:nth-of-type(1)",
+              rect: { x: 1, y: 2, width: 30, height: 40 },
+            },
+          ],
+          regions: [{ id: "region_2", rect: { x: 5, y: 6, width: 20, height: 30 } }],
+          strokes: [],
+          styleChanges: [],
+        };
+        const reload = (url: string) => {
+          listeners.get("did-start-navigation")?.({
+            url,
+            isSameDocument: false,
+            isMainFrame: true,
+            frame: null,
+          });
+        };
+
+        yield* manager.createTab("tab_1");
+        yield* manager.registerWebview("tab_1", 42);
+        const pick = yield* manager.pickElement("tab_1").pipe(Effect.forkChild);
+        yield* Effect.yieldNow;
+        onDraft?.({}, draft);
+        webviewSend.mockClear();
+
+        // A dev-server refresh reloads the page it is already on.
+        reload("https://example.com/app");
+        yield* Effect.yieldNow;
+        expect(pick.pollUnsafe()).toBeUndefined();
+
+        listeners.get("dom-ready")?.();
+        yield* Effect.yieldNow;
+        expect(pick.pollUnsafe()).toBeUndefined();
+        expect(webviewSend).toHaveBeenCalledWith("preview:start-pick", expect.anything(), draft);
+
+        // A reload that redirects away ends the pick, as other navigations do.
+        reload("https://example.com/app");
+        currentUrl = "https://example.com/login";
+        listeners.get("dom-ready")?.();
+        yield* Effect.yieldNow;
+        expect(yield* Fiber.join(pick)).toBeNull();
+      }),
+    ),
+  );
+
   effectIt.effect("settles the pick when the annotation screenshot never arrives", () =>
     withManager((manager) =>
       Effect.gen(function* () {

apps/desktop/src/preview/Manager.ts

@@ -72,6 +72,7 @@ import { PREVIEW_PICTURE_IN_PICTURE_FRAME_CHANNEL } from "../ipc/channels.ts";
 import * as BrowserSession from "./BrowserSession.ts";
 import {
   ANNOTATION_CAPTURED_CHANNEL,
+  ANNOTATION_DRAFT_CHANNEL,
   ANNOTATION_THEME_CHANNEL,
   CANCEL_PICK_CHANNEL,
   ELEMENT_PICKED_CHANNEL,
@@ -84,7 +85,11 @@ import {
   RECORDING_CONTROLLER_CHANNEL,
   START_PICK_CHANNEL,
 } from "./GuestProtocol.ts";
-import { isPreviewAnnotationPayload } from "./PickedElementPayload.ts";
+import {
+  isPreviewAnnotationDraft,
+  isPreviewAnnotationPayload,
+  type PreviewAnnotationDraft,
+} from "./PickedElementPayload.ts";
 import { playwrightInjectedRuntimeInstallExpression } from "./PlaywrightInjectedRuntime.ts";
 import {
   makePreviewAutomationKeySequence,
@@ -345,6 +350,10 @@ const normalizeCaptureRect = (value: unknown): PreviewAnnotationRect | null => {
   };
 };
 
+/** Whether two URLs load the same page, ignoring the fragment. */
+const isSamePageUrl = (left: string, right: string): boolean =>
+  left.split("#", 1)[0] === right.split("#", 1)[0];
+
 /** `capturePage` never settles when the guest's compositor is wedged. */
 const ANNOTATION_SCREENSHOT_TIMEOUT = "5 seconds";
 
@@ -2679,8 +2688,10 @@ const makeNativeOperations = Effect.fn("PreviewManager.makeOperations")(function
         const cleanup = Effect.fn("PreviewManager.cleanupPickElement")(function* () {
           yield* attempt({ operation: "pickElement.cleanup", tabId, webContentsId: wc.id }, () => {
             wc.ipc.removeListener(ELEMENT_PICKED_CHANNEL, onMessage);
+            wc.ipc.removeListener(ANNOTATION_DRAFT_CHANNEL, onDraft);
             wc.off("destroyed", onDestroyed);
             wc.off("did-start-navigation", onNavigated);
+            wc.off("dom-ready", onDomReady);
           }).pipe(Effect.ignore);
           // Only drop the slot while it is still ours. A newer session may
           // already have swapped itself in before cancelling this one.
@@ -2737,12 +2748,19 @@ const makeNativeOperations = Effect.fn("PreviewManager.makeOperations")(function
           }
           resume(Effect.succeed(null));
         });
+        // A reload of the page keeps the markup alive: the preload keeps the
+        // draft current here, and the fresh document gets it back on dom-ready.
+        const pageUrl = wc.getURL();
+        let draft: PreviewAnnotationDraft | null = null;
+        let reloading = false;
+        let submitted = false;
         const onMessage = (_event: Electron.IpcMainEvent, ...args: unknown[]): void => {
           const payload = args[0];
           if (!isPreviewAnnotationPayload(payload)) {
             settle(null);
             return;
           }
+          submitted = true;
           const cropRect = normalizeCaptureRect(args[1]);
           const submission = args[2] === "send" ? "send" : "attach";
           runFork(
@@ -2776,11 +2794,40 @@ const makeNativeOperations = Effect.fn("PreviewManager.makeOperations")(function
             ),
           );
         };
+        const onDraft = (_event: Electron.IpcMainEvent, value: unknown): void => {
+          if (!submitted && isPreviewAnnotationDraft(value)) draft = value;
+        };
         const onDestroyed = () => settle(null);
         const onNavigated = (
           event: Electron.Event<Electron.WebContentsDidStartNavigationEventParams>,
         ) => {
-          if (event.isMainFrame) settle(null);
+          if (!event.isMainFrame) return;
+          if (!event.isSameDocument && isSamePageUrl(event.url, pageUrl)) {
+            reloading = true;
+            return;
+          }
+          settle(null);
+        };
+        const onDomReady = () => {
+          if (!reloading) return;
+          reloading = false;
+          // A submitted annotation is already capturing and settles on its own.
+          if (submitted) return;
+          // The reload redirected somewhere else, so the markup no longer applies.
+          if (!isSamePageUrl(wc.getURL(), pageUrl)) {
+            settle(null);
+            return;
+          }
+          runFork(
+            Ref.get(annotationThemeRef).pipe(
+              Effect.flatMap((theme) =>
+                attempt({ operation: "pickElement.restore", tabId, webContentsId: wc.id }, () =>
+                  wc.send(START_PICK_CHANNEL, theme, draft),
+                ),
+              ),
+              Effect.ignore,
+            ),
+          );
         };
         const registerPickElement = Effect.fn("PreviewManager.registerPickElement")(function* () {
           // Two picks on one tab can overlap. Swap this session in and cancel
@@ -2800,8 +2847,10 @@ const makeNativeOperations = Effect.fn("PreviewManager.makeOperations")(function
           if (settled) return;
           yield* attempt({ operation: "pickElement.register", tabId, webContentsId: wc.id }, () => {
             wc.ipc.on(ELEMENT_PICKED_CHANNEL, onMessage);
+            wc.ipc.on(ANNOTATION_DRAFT_CHANNEL, onDraft);
             wc.once("destroyed", onDestroyed);
             wc.on("did-start-navigation", onNavigated);
+            wc.on("dom-ready", onDomReady);
             if (!wc.isFocused()) wc.focus();
             wc.send(START_PICK_CHANNEL, annotationTheme);
           });