platform: fix adding a second annotation (#441)

Author: irahopkinson

libs/shared-react/src/plugins/usj/annotation/AnnotationPlugin.test.tsx

@@ -2,9 +2,9 @@ import { baseTestEnvironment } from "../react-test.utils";
 import { AnnotationPlugin, AnnotationRef } from "./AnnotationPlugin";
 import { AnnotationRange } from "./selection.model";
 import { act } from "@testing-library/react";
-import { $createTextNode, $getRoot } from "lexical";
+import { $createTextNode, $getRoot, $isTextNode } from "lexical";
 import { createRef } from "react";
-import { $createParaNode, $createTypedMarkNode } from "shared";
+import { $createParaNode, $createTypedMarkNode, $isParaNode, $isTypedMarkNode } from "shared";
 import { vi } from "vitest";
 
 describe("AnnotationPlugin", () => {
@@ -30,6 +30,62 @@ describe("AnnotationPlugin", () => {
     // Check: spelling removal callback should not fire during internal rewrap.
     expect(spellingOnRemove).not.toHaveBeenCalled();
   });
+
+  it("correctly handles adding a second annotation at the same start position as the first", async () => {
+    // Initial state: "the " + "man" (spelling annotation) + " who stands"
+    const { annotationPlugin, editor } = await testEnvironment(() => {
+      $getRoot().append(
+        $createParaNode().append(
+          $createTextNode("the "),
+          $createTypedMarkNode({ spelling: ["spell-1"] }).append($createTextNode("man")),
+          $createTextNode(" who stands"),
+        ),
+      );
+    });
+    const jsonPath = "$.content[0].content[0]";
+    // Second annotation: "man who" (starts at same position as first but extends further)
+    const secondAnnotationRange: AnnotationRange = {
+      start: { jsonPath, offset: 4 },
+      end: { jsonPath, offset: 11 },
+    };
+
+    // SUT: add second annotation that starts at the same position as the first
+    await act(async () => {
+      annotationPlugin.setAnnotation(secondAnnotationRange, "grammar", "grammar-1");
+    });
+
+    // Verify the structure: both annotations should start at the same position
+    editor.getEditorState().read(() => {
+      const para = $getRoot().getFirstChild();
+      if (!$isParaNode(para)) throw new Error("Expected a ParaNode");
+
+      // First child should be plain text "the "
+      const firstChild = para.getFirstChild();
+      if (!$isTextNode(firstChild)) throw new Error("Expected a TextNode");
+      expect(firstChild.getTextContent()).toBe("the ");
+
+      // Second child should be a TypedMarkNode with both spelling AND grammar at "man"
+      const secondChild = firstChild.getNextSibling();
+      if (!$isTypedMarkNode(secondChild)) throw new Error("Expected a TypedMarkNode");
+      const secondTypedIDs = secondChild.getTypedIDs();
+      expect(secondTypedIDs.spelling).toContain("spell-1");
+      expect(secondTypedIDs.grammar).toContain("grammar-1");
+      expect(secondChild.getTextContent()).toBe("man");
+
+      // Third child should be a TypedMarkNode with only grammar for " who"
+      const thirdChild = secondChild.getNextSibling();
+      if (!$isTypedMarkNode(thirdChild)) throw new Error("Expected a TypedMarkNode");
+      const thirdTypedIDs = thirdChild.getTypedIDs();
+      expect(thirdTypedIDs.spelling).toBeUndefined();
+      expect(thirdTypedIDs.grammar).toContain("grammar-1");
+      expect(thirdChild.getTextContent()).toBe(" who");
+
+      // Fourth child should be plain text " stands"
+      const fourthChild = thirdChild.getNextSibling();
+      if (!$isTextNode(fourthChild)) throw new Error("Expected a TextNode");
+      expect(fourthChild.getTextContent()).toBe(" stands");
+    });
+  });
 });
 
 async function testEnvironment($initialEditorState: () => void) {

libs/shared-react/src/plugins/usj/annotation/selection.utils.ts

@@ -128,6 +128,11 @@ function $getNodeFromLocation(
       currentNode = $getContentChildAtIndex(currentNode, index);
     }
 
+    // If we landed on a TypedMarkNode (annotation wrapper), unwrap it to get the actual content
+    if (currentNode && $isTypedMarkNode(currentNode)) {
+      currentNode = currentNode.getFirstChild() ?? undefined;
+    }
+
     // If the jsonPath resolves to an ElementNode (e.g. "$.content[0]"), interpret offset as a
     // content-based child boundary offset (what the editor can emit), and return an element point.
     if (currentNode && $isElementNode(currentNode)) {
@@ -306,15 +311,20 @@ function $findTextNodeInMarks(
   if (!node || !$isTextNode(node)) return [undefined, undefined];
 
   const text = node.getTextContent();
-  if (offset >= 0 && offset <= text.length) return [node, offset];
+  // If offset is within the text (not at the very end), return this node.
+  // If offset equals text.length exactly, continue to next sibling to find the start of next content.
+  if (offset >= 0 && offset < text.length) return [node, offset];
 
   let nextNode = node.getNextSibling();
   if (!nextNode) {
     const parent = node.getParent();
     if ($isTypedMarkNode(parent)) nextNode = parent.getNextSibling();
   }
-  if (!nextNode || (!$isTypedMarkNode(nextNode) && !$isTextNode(nextNode)))
+  if (!nextNode || (!$isTypedMarkNode(nextNode) && !$isTextNode(nextNode))) {
+    // No more siblings - if offset equals text.length, return end of current node
+    if (offset === text.length) return [node, offset];
     return [undefined, undefined];
+  }
 
   const nextOffset = offset - text.length;
   if (nextNode && $isTextNode(nextNode)) return $findTextNodeInMarks(nextNode, nextOffset);