Keep highlights and drawings when saving immediately after reopening

scripts/smoke-test.mjs

@@ -463,8 +463,23 @@ try {
   await page.evaluate(() => document.querySelector('.page[data-page-number="1"]').scrollIntoView());
   await page.waitForTimeout(800);
   await page.keyboard.press("ControlOrMeta+w");
+  await page.getByRole("button", { name: /Try a sample/ }).waitFor();
+  // A slow notes read can finish after the reopened PDF is already usable. Hold the result until Save is pressed.
+  await page.evaluate(() => {
+    const transaction = IDBDatabase.prototype.transaction;
+    IDBDatabase.prototype.transaction = function(stores, mode, ...args) {
+      const tx = transaction.call(this, stores, mode, ...args);
+      if (mode !== "readonly" || !tx.objectStoreNames.contains("notes")) return tx;
+      IDBDatabase.prototype.transaction = transaction;
+      Object.defineProperty(tx, "oncomplete", { set(handler) {
+        tx.addEventListener("complete", event => { window.releaseSavedMarks = () => handler.call(tx, event); });
+      } });
+      return tx;
+    };
+  });
   await page.getByRole("button", { name: /Try a sample/ }).click();
   await page.locator(".textLayer span:not(.markedContent)").first().waitFor({ timeout: 30_000 });
+  await page.waitForFunction(() => typeof window.releaseSavedMarks === "function");
   const copy = path.join(env.FOLIO_TEST_USER_DATA, "copy.pdf");
   await app.evaluate(({ session }, file) => {
     globalThis.folioCopySaved = new Promise((resolve) => session.defaultSession.once("will-download", (_event, item) => {
@@ -473,6 +488,7 @@ try {
     }));
   }, copy);
   await page.keyboard.press("ControlOrMeta+s");
+  await page.evaluate(() => window.releaseSavedMarks());
   assert.equal(await app.evaluate(() => globalThis.folioCopySaved), "completed");
   assert.equal(readFileSync(copy, "latin1").match(/\/Subtype\s*\/Highlight\b/g)?.length, 4);
   const savedInk = await page.evaluate(async bytes => {

scripts/test-reader-security.mjs

@@ -176,16 +176,17 @@ try {
     const mark = { id: "saved", type: "highlight", page: 0, text: "saved selection", note: "", color: "#00ff00",
       createdAt: 1, data: { color: [0, 255, 0] } };
     await saveNotes({ id: "reopened", marks: [mark], note: "", updatedAt: 1 }, 0);
+    const started = Promise.withResolvers();
     const deferred = Promise.withResolvers();
     let staleAdds = 0, currentAdds = 0;
     const firstLoad = markup.load("reopened", true);
-    manager({ div: root, deserialize: () => deferred.promise, add() { staleAdds++; } });
-    await firstLoad;
+    manager({ div: root, deserialize() { started.resolve(); return deferred.promise; }, add() { staleAdds++; } });
+    await started.promise;
     const nextLoad = markup.load("reopened", true);
     manager({ div: root, deserialize: async () => editor(), add() { currentAdds++; } });
     await nextLoad;
     deferred.resolve(editor());
-    await new Promise(resolve => setTimeout(resolve, 0));
+    await firstLoad;
     markup.unload();
     root.remove();
     return { firstText, highlightedText, staleAdds, currentAdds };

src/App.tsx

@@ -98,6 +98,7 @@ interface OpenDoc {
   /** The file's size; its bytes live in pdf.js's worker (pdf.getData()) rather than in a second copy here. */
   size: number;
   pdf: PDFDocumentProxy;
+  markupReady: Promise<void>;
   /** Used only in memory when reopening an annotated snapshot for page reordering. */
   password?: string;
   /** Original library entry while content edits are held in memory. */
@@ -439,16 +440,16 @@ export default function App({ tabs, tabKey, upgradeNotes }: { tabs: DocumentTabs
         }
 
         const prev = docRef.current;
-        const next: OpenDoc = { id, name, size: data.byteLength, pdf, password, sourceId: revision?.sourceId ?? revision?.id, historyAllowed, edits };
-        docRef.current = next;
         // Password-protected documents aren't stored by Folio (same rule as Recent).
         const reading = previous ? markupRef.current!.state : undefined;
         if (reading?.bookmarkedPages && exportedOrder) {
           const bookmarks = new Set(reading.bookmarkedPages);
           reading.bookmarkedPages = exportedOrder.flatMap((n, i) => bookmarks.has(n) ? [i + 1] : []);
         }
         persistRef.current = !password && !revision;
-        void markupRef.current!.load(id, persistRef.current, reading);
+        const markupReady = markupRef.current!.load(id, persistRef.current, reading);
+        const next: OpenDoc = { id, name, size: data.byteLength, pdf, password, markupReady, sourceId: revision?.sourceId ?? revision?.id, historyAllowed, edits };
+        docRef.current = next;
         setNoteBox(null);
         setPinned(null);
         setAiContext(null);
@@ -867,6 +868,8 @@ export default function App({ tabs, tabKey, upgradeNotes }: { tabs: DocumentTabs
     const storage = doc.pdf.annotationStorage;
     const keys: string[] = [];
     try {
+      await doc.markupReady;
+      if (docRef.current !== doc) return null;
       const unrendered = await markupRef.current!.unrenderedMarks(doc.pdf);
       for (const data of unrendered) {
         const key = `pdfjs_internal_editor_folio-unrendered-${keys.length}`;
@@ -904,8 +907,10 @@ export default function App({ tabs, tabKey, upgradeNotes }: { tabs: DocumentTabs
   const exportNotes = async (format: "markdown" | "json") => {
     if (!doc) return;
     const controller = markupRef.current!;
-    controller.sync();
     try {
+      await doc.markupReady;
+      if (docRef.current !== doc) return;
+      controller.sync();
       if (format === "markdown") {
         const ok = await copyText(notesMarkdown(doc.name, controller.state));
         if (docRef.current === doc) flash(ok ? "Markdown copied" : "Copy failed");

src/markup.ts

@@ -167,8 +167,8 @@ export class MarkupController {
     }
     this.#emit();
     if (reading) this.#scheduleSave();
-    // Pages that rendered before the saved data arrived.
-    for (const page of [...this.#pending.keys()]) if (this.#ui?.getLayer(page)) void this.#restorePage(page);
+    // Finish restoring visible pages before an export collects the marks still pending on other pages.
+    await Promise.allSettled(Array.from(this.#pending.keys(), page => this.#restorePage(page)));
   }
 
   unload() {