Fix CKEditor 4 typing after frame rewrites, a watchdog restart, and a lost arrow-key choice

src/adapters/chrome/content-script/HostChangeWatcher.ts

@@ -68,7 +68,15 @@ export class HostChangeWatcher {
     }
 
     const currentNode = document.body || document.documentElement;
-    if (this.dependencies.getObservedNode() === currentNode) {
+    const observedNode = this.dependencies.getObservedNode();
+    if (observedNode === currentNode) {
+      return;
+    }
+    // A script that started while the page was loading, before the body (CKEditor 4
+    // writes its editing frame so), observes the root. The root still holds the body:
+    // a restart would only drop the open suggestions.
+    if (observedNode.isConnected && observedNode.contains(currentNode)) {
+      this.dependencies.setObservedNode(currentNode);
       return;
     }
 

src/adapters/chrome/content-script/content_script.ts

@@ -48,6 +48,17 @@ class FluentTyper {
       this.runtimeController.handleEarlyTabAcceptRequest(event.data.entryId);
     }
   };
+  // document.open() (CKEditor 4 writes its editing frame so, also on setData) erases
+  // every listener of this script and replaces the root element. Only an observer of
+  // the document itself sees that: then a new instance starts on the new document.
+  private readonly root = document.documentElement;
+  private readonly documentRewriteObserver = new MutationObserver(() => {
+    if (!document.documentElement || document.documentElement === this.root) return;
+    logger.info("Document rewritten; restarting content script");
+    this.destroy();
+    window.FluentTyper = new FluentTyper();
+  });
+  private destroyed = false;
 
   constructor() {
     logger.info("Initializing content script", {
@@ -88,6 +99,7 @@ class FluentTyper {
     });
 
     chrome.runtime.onMessage.addListener(this.boundMessageHandler);
+    this.documentRewriteObserver.observe(document, { childList: true });
     this.getConfig();
   }
 
@@ -137,6 +149,8 @@ class FluentTyper {
 
   destroy(): void {
     logger.info("Destroying content script instance");
+    this.destroyed = true;
+    this.documentRewriteObserver.disconnect();
     this.hostChangeWatcher.stop();
     this.runtimeController.disable();
     window.removeEventListener("message", this.boundEarlyTabAcceptHandler);
@@ -154,7 +168,8 @@ class FluentTyper {
     };
     chrome.runtime.sendMessage(msg, (response: unknown) => {
       checkLastError();
-      this.messageHandler(response as Message);
+      // A late answer must not enable a destroyed instance again.
+      if (!this.destroyed) this.messageHandler(response as Message);
     });
   }
 }

src/adapters/chrome/content-script/suggestions/SuggestionEntrySession.ts

@@ -304,6 +304,8 @@ export class SuggestionEntrySession {
       this.deferredInput = event;
       return;
     }
+    // An edit makes a new choice necessary.
+    this.entry.chosenSuggestion = null;
     if (!this.refreshInteraction()) {
       if (this.entry.isComposing) this.handleSuppressedInput();
       return;
@@ -404,6 +406,7 @@ export class SuggestionEntrySession {
     this.snippetSuggestions.clear();
     this.snippetShortcuts = undefined;
     this.entry.selectedIndex = 0;
+    this.entry.chosenSuggestion = null;
     this.entry.visibleSuggestionBeforeCursorText = null;
     this.entry.visibleSuggestionFullText = null;
     this.entry.inlineSuggestion = null;
@@ -442,7 +445,12 @@ export class SuggestionEntrySession {
     this.snippetSuggestions = new Set(
       this.entry.suggestions.filter((_, index) => context.snippetShortcuts?.[index]),
     );
-    this.entry.selectedIndex = 0;
+    // An answer that comes after an arrow key (it was late) keeps the user's choice.
+    const chosen = this.entry.chosenSuggestion
+      ? this.entry.suggestions.indexOf(this.entry.chosenSuggestion)
+      : -1;
+    if (chosen < 0) this.entry.chosenSuggestion = null;
+    this.entry.selectedIndex = Math.max(0, chosen);
     this.entry.menuHeader =
       this.options.showSuggestionFooter && context.lang && SUPPORTED_LANGUAGES[context.lang]
         ? suggestionLanguageLabel(SUPPORTED_LANGUAGES[context.lang])

src/adapters/chrome/content-script/suggestions/SuggestionKeyboardHandler.ts

@@ -154,6 +154,7 @@ export class SuggestionKeyboardHandler {
     if (next < rows) {
       entry.selectedIndex = next;
     }
+    entry.chosenSuggestion = next < rows ? entry.suggestions[next] : null;
     this.options.updateSelectionHighlight(entry);
   }
 

src/adapters/chrome/content-script/suggestions/SuggestionManagerRuntime.ts

@@ -572,6 +572,7 @@ export class SuggestionManagerRuntime {
       requestId: 0,
       suggestions: [],
       selectedIndex: 0,
+      chosenSuggestion: null,
       menuHeader: null,
       latestMentionText: "",
       latestMentionStart: 0,

src/adapters/chrome/content-script/suggestions/types.ts

@@ -134,6 +134,8 @@ export interface SuggestionEntry {
   requestId: number;
   suggestions: string[];
   selectedIndex: number;
+  /** The suggestion the user moved to with an arrow key after the last edit. A later answer keeps it highlighted. */
+  chosenSuggestion: string | null;
   menuHeader: string | null;
   latestMentionText: string;
   latestMentionStart: number;

tests/SuggestionEntrySession.test.ts

@@ -1161,6 +1161,41 @@ test("session ignores stale responses and renders fresh menu responses", () => {
   expect(recordSuggestionShown).toHaveBeenCalledWith({ suggestionCount: 1, language: "en_US" });
 });
 
+test("a late answer keeps the suggestion that the user chose with an arrow key", () => {
+  const renderMenu = jest.fn();
+  const entry = createSuggestionEntry({ requestId: 2, latestMentionText: "th" });
+  const input = entry.elem as HTMLInputElement;
+  input.value = "th";
+  input.selectionStart = 2;
+  input.selectionEnd = 2;
+  const session = makeSession({ entry, renderMenu });
+  const answer = (predictions: string[]) =>
+    session.handlePredictionResponse(
+      partialResponse({ requestId: 2, suggestionId: 1, predictions }),
+    );
+
+  answer(["the", "that", "this"]);
+  // ArrowDown: the keyboard handler records the choice.
+  entry.selectedIndex = 1;
+  entry.chosenSuggestion = "that";
+  // The answer for the typed text comes after the arrow key, in another order.
+  answer(["this", "the", "that"]);
+  expect(entry.selectedIndex).toBe(2);
+  expect(renderMenu).toHaveBeenLastCalledWith(expect.objectContaining({ selectedIndex: 2 }));
+
+  // An answer without the chosen suggestion highlights the first row again.
+  answer(["they", "them"]);
+  expect(entry.selectedIndex).toBe(0);
+  expect(entry.chosenSuggestion).toBeNull();
+
+  // Typing makes a new choice necessary.
+  entry.chosenSuggestion = "them";
+  input.value = "the";
+  input.selectionStart = input.selectionEnd = 3;
+  session.handleInput(new Event("input"));
+  expect(entry.chosenSuggestion).toBeNull();
+});
+
 test("session does not fulfill pending inline accept when the ghost render is vetoed", () => {
   const textEditService = {
     acceptSuggestion: jest.fn(() => null),

tests/SuggestionKeyboardHandler.test.ts

@@ -32,6 +32,20 @@ describe("SuggestionKeyboardHandler", () => {
     expect(updateSelectionHighlight).toHaveBeenCalledWith(entry);
   });
 
+  test("arrow keys record the chosen suggestion for a later answer", () => {
+    const handler = createHandler({
+      inlineSuggestionEnabled: false,
+      consumeKeyboardEvent: jest.fn((event: KeyboardEvent) => event.preventDefault()),
+      isMenuVisible: jest.fn(() => true),
+    });
+    const entry = createSuggestionEntry({ suggestions: ["one", "two"], selectedIndex: 0 });
+
+    handler.handle(entry, createEvent("ArrowDown"));
+    expect(entry.chosenSuggestion).toBe("two");
+    handler.handle(entry, createEvent("ArrowDown"));
+    expect(entry.chosenSuggestion).toBe("one");
+  });
+
   test("moves selection the other way when the menu lists suggestions bottom-up", () => {
     const handler = createHandler({
       inlineSuggestionEnabled: false,

tests/content_script.behavior.test.ts

@@ -18,6 +18,10 @@ import {
   EARLY_TAB_ACCEPT_MESSAGE_TYPE,
   EARLY_TAB_ACCEPT_REQUEST_EVENT,
 } from "../src/adapters/chrome/content-script/suggestions/EarlyTabAcceptBridgeProtocol";
+import {
+  HOST_EDITOR_ENABLED_ATTR,
+  HOST_EDITOR_ENABLED_EVENT,
+} from "../src/adapters/chrome/content-script/suggestions/HostEditorBridgeProtocol";
 
 type SuggestionLike = {
   queryAndAttachHelper: jest.Mock;
@@ -882,6 +886,20 @@ describe("content_script behavior", () => {
     expect(sendMessage).toHaveBeenCalled();
   });
 
+  test("watchdog observes a body that appeared under the observed root without a restart", async () => {
+    const { fluentTyper, domObserverInstances } = await loadContentScript();
+    const domObserver = domObserverInstances[0];
+    const restartSpy = jest.spyOn(fluentTyper, "restart");
+    fluentTyper.enabled = true;
+    // The script started while the page was loading: there was no body, so it observes the root.
+    domObserver.getNode.mockReturnValue(document.documentElement);
+
+    fluentTyper.watchDog();
+
+    expect(restartSpy).not.toHaveBeenCalled();
+    expect(domObserver.setNode).toHaveBeenCalledWith(document.body);
+  });
+
   test("watchdog prefers a host change over a node change and skips the DOM restart", async () => {
     const { fluentTyper, domObserverInstances } = await loadContentScript();
     const domObserver = domObserverInstances[0];
@@ -898,6 +916,46 @@ describe("content_script behavior", () => {
     expect(domObserver.setNode).not.toHaveBeenCalled();
   });
 
+  test("a rewritten document (document.open) starts a new instance that enables the host bridge again", async () => {
+    const { fluentTyper, sendMessage } = await loadContentScript();
+    const lateConfig = sendMessage.mock.calls[0][1] as (response: unknown) => void;
+    fluentTyper.setConfig(defaultConfig());
+    expect(fluentTyper.enabled).toBe(true);
+
+    // document.open() erases every listener and replaces the root element.
+    const oldRoot = document.documentElement;
+    const newRoot = document.createElement("html");
+    newRoot.append(document.createElement("head"), document.createElement("body"));
+    document.replaceChild(newRoot, oldRoot);
+    try {
+      await Promise.resolve();
+      const restarted = (window as Window & { FluentTyper?: LoadedContentScript["fluentTyper"] })
+        .FluentTyper!;
+      behaviorHarness.fluentTyperInstances.push(restarted);
+      expect(restarted).not.toBe(fluentTyper);
+      expect(sendMessage.mock.calls.at(-1)?.[0]).toEqual(
+        expect.objectContaining({ command: CMD_CONTENT_SCRIPT_GET_CONFIG }),
+      );
+
+      // A late config answer to the old instance does not enable it again.
+      const handled = jest.spyOn(fluentTyper, "messageHandler");
+      lateConfig({ command: CMD_BACKGROUND_PAGE_SET_CONFIG, context: defaultConfig() });
+      expect(handled).not.toHaveBeenCalled();
+
+      // The main-world bridge lost its state: the new instance sends "enabled" again.
+      const bridgeStates: (string | null)[] = [];
+      const record = () =>
+        bridgeStates.push(document.documentElement.getAttribute(HOST_EDITOR_ENABLED_ATTR));
+      document.addEventListener(HOST_EDITOR_ENABLED_EVENT, record);
+      restarted.setConfig(defaultConfig());
+      document.removeEventListener(HOST_EDITOR_ENABLED_EVENT, record);
+      expect(bridgeStates).toEqual(["true"]);
+    } finally {
+      for (const instance of behaviorHarness.fluentTyperInstances) instance.destroy();
+      document.replaceChild(oldRoot, newRoot);
+    }
+  });
+
   test("watchdog does nothing when host and observed node are unchanged", async () => {
     const { fluentTyper, domObserverInstances } = await loadContentScript();
     const domObserver = domObserverInstances[0];

tests/e2e/coverage-matrix.json

@@ -4703,6 +4703,38 @@
         }
       ]
     },
+    {
+      "id": "rewritten_document_restart",
+      "description": "When a page rewrites a document with document.open() (CKEditor 4 writes its editing frame so, also on setData), every listener in it is erased and its root element is replaced. FluentTyper then starts a new content script instance on the new document, which sends the enabled state to the main-world bridge again; a late config answer does not enable the old instance. Typing in the rewritten CKEditor 4 frame accepts a prediction with Tab.",
+      "coverage": [
+        {
+          "layer": "unit",
+          "file": "tests/content_script.behavior.test.ts",
+          "test": "a rewritten document (document.open) starts a new instance that enables the host bridge again"
+        },
+        {
+          "layer": "e2e-full",
+          "file": "tests/e2e/full.e2e.test.ts",
+          "test": "CKEditor 4 setData rewrites its frame with document.open, and typing still accepts a prediction"
+        }
+      ]
+    },
+    {
+      "id": "late_body_without_restart",
+      "description": "A content script that starts while its document is still loading, before the body (CKEditor 4 writes its editing frame so), observes the root element. When the body appears under that root, the watchdog observes the body and does not restart the runtime, so a visible menu stays and Tab still accepts from it.",
+      "coverage": [
+        {
+          "layer": "unit",
+          "file": "tests/content_script.behavior.test.ts",
+          "test": "watchdog observes a body that appeared under the observed root without a restart"
+        },
+        {
+          "layer": "e2e-full",
+          "file": "tests/e2e/full.e2e.test.ts",
+          "test": "%s typing expands a snippet in one native undo step"
+        }
+      ]
+    },
     {
       "id": "rich_editor_snippet_expansion",
       "description": "In real Draft.js, Trix, CKEditor 4, Froala, Summernote, Tiptap and RoosterJS editors and in bundled Quill 1 and Quill 2, a text expansion replaces its shortcut in the second paragraph of the editor's own model, keeps bold and links, and native Undo restores the shortcut in one step. A bundled Quill 2's Undo can also remove the typing of the last second.",
@@ -4714,6 +4746,27 @@
         }
       ]
     },
+    {
+      "id": "late_answer_keeps_arrow_choice",
+      "description": "When the user moves the menu selection with an arrow key and a prediction answer comes after that (it was late), the menu keeps the chosen suggestion highlighted and Tab accepts it. When the new answer does not contain the chosen suggestion, the first row is highlighted. An edit clears the choice.",
+      "coverage": [
+        {
+          "layer": "unit",
+          "file": "tests/SuggestionKeyboardHandler.test.ts",
+          "test": "arrow keys record the chosen suggestion for a later answer"
+        },
+        {
+          "layer": "unit",
+          "file": "tests/SuggestionEntrySession.test.ts",
+          "test": "a late answer keeps the suggestion that the user chose with an arrow key"
+        },
+        {
+          "layer": "e2e-full",
+          "file": "tests/e2e/full.e2e.test.ts",
+          "test": "A late prediction answer keeps the suggestion that the user chose with an arrow key"
+        }
+      ]
+    },
     {
       "id": "stale_predictions_after_caret_move",
       "description": "Predictions are for one caret place. When the page moves the caret without a text change (for example after focus), a prediction response for the old place is not shown, and a visible contenteditable menu is dismissed, so Space or Tab cannot write a stale word at the new place.",