fix: preserve AI line highlight against ProseMirror DOMObserver and prevent crash on accepting suggestions

2026-10-30 12:20:13 -05:00
parent 271d1e2c79
commit 2940e3c9aa
6 changed files with 240 additions and 39 deletions
@@ -80,6 +80,7 @@ export function WritingView({ onChapterChange }: WritingViewProps) {
  const autosave = useProjectStore((s) => s.autosave);
  const entities = useProjectStore((s) => s.entities);
  const highlightedLine = useProjectStore((s) => s.highlightedLine);
+  const highlightSeq = useProjectStore((s) => s.highlightSeq);
  const docExternalRev = useProjectStore((s) => s.docExternalRev);

  const fontFamily = useSettings((st) => st.fontFamily);
@@ -134,32 +135,29 @@ export function WritingView({ onChapterChange }: WritingViewProps) {
  // Sync content when the active chapter changes or when an AI suggestion edit is applied.
  useEffect(() => {
    if (!editor || !chapter) return;
-    editor.commands.setContent(chapter.doc, { emitUpdate: false });
-    editor.commands.focus();
+    try {
+      editor.commands.setContent(chapter.doc, { emitUpdate: false });
+    } catch {
+      /* ignore invalid intermediate content */
+    }
  }, [editor, chapter?.id, docExternalRev]);

-  // Highlight a specific 1-based line block in the editor when an AI suggestion targets it.
+  // Scroll the highlighted 1-based line block into view when an AI suggestion targets it.
+  // Visual highlighting uses CSS `:nth-child(N)` via `data-highlight-line` on `.ps-editor`
+  // so ProseMirror's internal DOMObserver never strips the highlight class.
  useEffect(() => {
    const editorDom = editorDomRef.current;
-    if (!editorDom) return undefined;
+    if (!editorDom || highlightedLine == null || highlightedLine < 1) return undefined;
    const pmRoot = editorDom.querySelector(".ProseMirror");
    if (!pmRoot) return undefined;

    const children = Array.from(pmRoot.children) as HTMLElement[];
-    children.forEach((el) => el.classList.remove("ps-ai-line-highlight"));
-
-    if (highlightedLine != null && highlightedLine >= 1 && highlightedLine <= children.length) {
-      const targetEl = children[highlightedLine - 1];
-      if (targetEl) {
-        targetEl.classList.add("ps-ai-line-highlight");
-        targetEl.scrollIntoView({ behavior: "smooth", block: "center" });
-      }
+    const targetEl = children[highlightedLine - 1];
+    if (targetEl && typeof targetEl.scrollIntoView === "function") {
+      targetEl.scrollIntoView({ behavior: "smooth", block: "center" });
    }
-
-    return () => {
-      children.forEach((el) => el.classList.remove("ps-ai-line-highlight"));
-    };
-  }, [highlightedLine, docExternalRev, chapter?.id]);
+    return undefined;
+  }, [editor, highlightedLine, highlightSeq, docExternalRev, chapter?.id]);

  // Focus mode: dim every top-level block in .ProseMirror except the one
  // containing the active cursor selection.
@@ -859,7 +857,28 @@ export function WritingView({ onChapterChange }: WritingViewProps) {
          className={`ps-editor-scroll${focusMode ? " ps-editor-focus" : ""}${showLineNumbers || highlightedLine != null ? " ps-editor-linenumbers" : ""}${firstLineIndent ? " ps-editor-indent" : ""}`}
        >
          <div className="ps-editor-writing">
-            <div className="ps-editor" ref={editorDomRef}>
+            {highlightedLine != null && highlightedLine >= 1 && (
+              <style>{`
+                .ps-editor[data-highlight-line="${highlightedLine}"] .ProseMirror > *:nth-child(${highlightedLine}) {
+                  background: color-mix(in srgb, var(--ps-accent) 24%, transparent) !important;
+                  outline: 2px solid var(--ps-accent) !important;
+                  outline-offset: 4px !important;
+                  border-radius: 6px !important;
+                  box-shadow: 0 0 0 4px color-mix(in srgb, var(--ps-accent) 14%, transparent) !important;
+                  transition: background 0.18s ease, outline 0.18s ease;
+                }
+                .ps-editor-linenumbers .ps-editor[data-highlight-line="${highlightedLine}"] .ProseMirror > *:nth-child(${highlightedLine})::before {
+                  color: var(--ps-accent) !important;
+                  font-weight: 800 !important;
+                  opacity: 1 !important;
+                }
+              `}</style>
+            )}
+            <div
+              className="ps-editor"
+              ref={editorDomRef}
+              data-highlight-line={highlightedLine ?? undefined}
+            >
              <EditorContent editor={editor} />
              <MentionPopover
                entities={entities}
@@ -1,6 +1,11 @@
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { useNavigate } from "react-router-dom";
-import { AIStore, AgentController, createIndexedDBPersistence } from "@pubscript/ai";
+import {
+  AIStore,
+  AgentController,
+  createIndexedDBPersistence,
+  createInMemoryPersistence,
+} from "@pubscript/ai";
 import {
  countTokens,
  type AIMessage,
@@ -366,7 +371,11 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {

  useEffect(() => {
    if (!projectId) return undefined;
-    const store = new AIStore(projectId, createIndexedDBPersistence());
+    const persistence =
+      typeof indexedDB !== "undefined"
+        ? createIndexedDBPersistence()
+        : createInMemoryPersistence();
+    const store = new AIStore(projectId, persistence);
    setAiStore(store);
    return () => {
      void store.dispose();
@@ -440,7 +449,7 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {
          lineNumber: adjusted.lineNumber,
          ...(adjusted.originalText ? { originalText: adjusted.originalText } : {}),
        };
-        if (aiStore && activeSessionId) {
+        if (aiStore && activeSessionId && typeof aiStore.updateSuggestion === "function") {
          void aiStore.updateSuggestion(activeSessionId, sg.id, {
            lineNumber: adjusted.lineNumber,
            ...(adjusted.originalText ? { originalText: adjusted.originalText } : {}),
@@ -461,7 +470,7 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {
      chapterId: currentChapterId,
      lines: chapterLines,
    };
-  }, [chapterLines, activeChapter?.id, suggestions, aiStore, activeSessionId, highlightedLine, setHighlightedLine]);
+  }, [chapterLines, activeChapter?.id]);

  // Sidebar drag-resize handler with snap-to-default (370px)
  const handleStartResize = (e: React.MouseEvent<HTMLDivElement>) => {
@@ -632,10 +641,12 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {
  }, [reloadAll]);

  useEffect(() => {
-    scrollRef.current?.scrollTo({
-      top: scrollRef.current.scrollHeight,
-      behavior: "smooth",
-    });
+    if (scrollRef.current && typeof scrollRef.current.scrollTo === "function") {
+      scrollRef.current.scrollTo({
+        top: scrollRef.current.scrollHeight,
+        behavior: "smooth",
+      });
+    }
  }, [messages, sending, thinkingStage]);

  const handleLoadReload = async () => {
@@ -894,6 +905,15 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {
    const replacement = (draftEdits[sg.id] ?? sg.newText ?? sg.text).trim();
    if (!replacement) return;

+    // Remove the accepted suggestion from state and persistence first so line
+    // re-anchoring only runs on the remaining pending suggestions.
+    setSuggestions((prev) => prev.filter((item) => item.id !== sg.id));
+    if (typeof aiStore.deleteSuggestion === "function") {
+      await aiStore.deleteSuggestion(activeSessionId, sg.id);
+    } else {
+      await aiStore.resolveSuggestion(activeSessionId, sg.id, true);
+    }
+
    const applied = await applyLineEdit({
      ...(sg.chapterId ? { chapterId: sg.chapterId } : {}),
      ...(sg.lineNumber != null ? { lineNumber: sg.lineNumber } : {}),
@@ -901,7 +921,6 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {
      newText: replacement,
    });

-    await aiStore.deleteSuggestion(activeSessionId, sg.id);
    const nextList = await aiStore.listSuggestions(activeSessionId);
    setSuggestions(nextList);

@@ -914,7 +933,12 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {

  const handleDenySuggestion = async (sg: AISuggestion) => {
    if (!aiStore || !activeSessionId) return;
-    await aiStore.deleteSuggestion(activeSessionId, sg.id);
+    setSuggestions((prev) => prev.filter((item) => item.id !== sg.id));
+    if (typeof aiStore.deleteSuggestion === "function") {
+      await aiStore.deleteSuggestion(activeSessionId, sg.id);
+    } else {
+      await aiStore.resolveSuggestion(activeSessionId, sg.id, false);
+    }
    const nextList = await aiStore.listSuggestions(activeSessionId);
    setSuggestions(nextList);
    if (highlightedLine === sg.lineNumber) {
@@ -960,11 +984,7 @@ export function AISidebar({ fullPage = false }: AISidebarProps) {

  const handleHighlightLine = (lineNum: number | undefined) => {
    if (!lineNum) return;
-    if (highlightedLine === lineNum) {
-      setHighlightedLine(null);
-    } else {
-      setHighlightedLine(lineNum);
-    }
+    setHighlightedLine(lineNum);
  };

  const renderMessageWithLineLinks = (content: string) => {
@@ -0,0 +1,136 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { createRoot } from "react-dom/client";
+import { act } from "react";
+import { MemoryRouter } from "react-router-dom";
+import type { Chapter, ProjectMeta, StorageAdapter } from "@pubscript/core";
+import { useProjectStore } from "../../store/project-store.js";
+import { WritingView } from "../../components/WritingView.js";
+import { AISidebar } from "./AISidebar.js";
+
+describe("AISidebar + WritingView accept suggestion flow", () => {
+  beforeEach(() => {
+    (globalThis as unknown as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
+    Element.prototype.scrollIntoView = () => {};
+    Range.prototype.getClientRects = () =>
+      ({
+        length: 0,
+        item: () => null,
+        [Symbol.iterator]: function* () {},
+      }) as unknown as DOMRectList;
+    Range.prototype.getBoundingClientRect = () =>
+      ({ top: 0, left: 0, bottom: 0, right: 0, width: 0, height: 0, x: 0, y: 0, toJSON: () => ({}) }) as DOMRect;
+    vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("offline"));
+
+    const ch: Chapter = {
+      id: "ch-1",
+      bookId: "book-1",
+      projectId: "proj-1",
+      title: "Chapter One",
+      order: 0,
+      rev: 1,
+      meta: { wordCount: 10, scenes: [], updatedAt: new Date().toISOString() },
+      doc: {
+        type: "doc",
+        content: [
+          {
+            type: "paragraph",
+            content: [
+              {
+                type: "text",
+                text: "The ferry that carried Mara Vell to the Glass Harbor left the mainland before dawn.",
+              },
+            ],
+          },
+          {
+            type: "paragraph",
+            content: [
+              {
+                type: "text",
+                text: "She pressed her palm against the cold window and watched the lighthouse appear out of the mist.",
+              },
+            ],
+          },
+          {
+            type: "paragraph",
+            content: [
+              {
+                type: "text",
+                text: "Three years since she had last seen it, and yet it seemed to have shrunk.",
+              },
+            ],
+          },
+        ],
+      },
+    };
+    const mockAdapter = {
+      async saveDocument() {},
+    } as unknown as StorageAdapter;
+
+    useProjectStore.setState({
+      adapter: mockAdapter,
+      project: { id: "proj-1", title: "Test Project" } as ProjectMeta,
+      chapters: [ch],
+      activeChapterId: "ch-1",
+      highlightedLine: null,
+      docExternalRev: 0,
+      status: "ready",
+    });
+  });
+
+  it("generates line suggestions, highlights line on click, and accepts one without crashing", async () => {
+    const container = document.createElement("div");
+    document.body.appendChild(container);
+    const root = createRoot(container);
+
+    await act(async () => {
+      root.render(
+        <MemoryRouter>
+          <div style={{ display: "flex" }}>
+            <WritingView />
+            <AISidebar />
+          </div>
+        </MemoryRouter>,
+      );
+    });
+
+    await act(async () => {
+      await new Promise((r) => setTimeout(r, 80));
+    });
+
+    const genBtn = Array.from(container.querySelectorAll("button")).find((b) =>
+      b.textContent?.includes("Suggest Line Edits for Chapter"),
+    );
+    expect(genBtn).toBeDefined();
+
+    await act(async () => {
+      genBtn!.click();
+    });
+    await act(async () => {
+      await new Promise((r) => setTimeout(r, 1600));
+    });
+
+    const acceptBtn = Array.from(container.querySelectorAll("button")).find((b) =>
+      b.textContent?.includes("Accept"),
+    );
+    if (!acceptBtn) {
+      console.log("BUTTONS:", Array.from(container.querySelectorAll("button")).map((b) => b.textContent));
+    }
+    expect(acceptBtn).toBeDefined();
+
+    await act(async () => {
+      acceptBtn!.click();
+      await new Promise((r) => setTimeout(r, 150));
+    });
+
+    const pm = container.querySelector(".ProseMirror");
+    expect(pm?.textContent).toContain("sharp and unmistakable");
+
+    await act(async () => {
+      root.unmount();
+    });
+    container.remove();
+  });
+});
@@ -70,6 +70,7 @@ interface ProjectStore {
  chapters: Chapter[];
  activeChapterId: string | null;
  highlightedLine: number | null;
+  highlightSeq: number;
  docExternalRev: number;
  status: "idle" | "loading" | "ready";
  error: string | null;
@@ -204,6 +205,7 @@ export const useProjectStore = create<ProjectStore>((set, get) => {
    chapters: [],
    activeChapterId: null,
    highlightedLine: null,
+    highlightSeq: 0,
    docExternalRev: 0,
    status: "idle",
    error: null,
@@ -538,7 +540,10 @@ export const useProjectStore = create<ProjectStore>((set, get) => {
    },

    setHighlightedLine: (line) => {
-      set({ highlightedLine: line });
+      set((state) => ({
+        highlightedLine: line,
+        highlightSeq: state.highlightSeq + 1,
+      }));
    },

    applyLineEdit: async ({ chapterId, lineNumber, originalText, newText }) => {
@@ -591,18 +596,28 @@ export const useProjectStore = create<ProjectStore>((set, get) => {

      const oldBlock = blocks[targetIndex]! as unknown as Record<string, unknown>;
      const oldFullText = extractBlockText(oldBlock);
-      const updatedLineText =
+      const updatedLineText = (
        originalText &&
        originalText.trim() &&
        oldFullText.includes(originalText.trim()) &&
        oldFullText.trim() !== originalText.trim()
          ? oldFullText.replace(originalText.trim(), newText)
-          : newText;
+          : newText
+      ).trim();
+
+      const safeType =
+        oldBlock.type === "heading" || oldBlock.type === "codeBlock"
+          ? (oldBlock.type as string)
+          : "paragraph";

      blocks[targetIndex] = {
-        ...blocks[targetIndex]!,
-        type: (oldBlock.type as string) || "paragraph",
-        content: [{ type: "text", text: updatedLineText }],
+        type: safeType,
+        ...(oldBlock.attrs && typeof oldBlock.attrs === "object"
+          ? { attrs: oldBlock.attrs as Record<string, unknown> }
+          : {}),
+        ...(updatedLineText
+          ? { content: [{ type: "text", text: updatedLineText }] }
+          : {}),
      };

      const nextDoc = {
@@ -622,6 +637,7 @@ export const useProjectStore = create<ProjectStore>((set, get) => {
      set((state) => ({
        activeChapterId: targetChapter.id,
        highlightedLine: targetIndex + 1,
+        highlightSeq: state.highlightSeq + 1,
        docExternalRev: state.docExternalRev + 1,
        chapters: state.chapters.map((c) =>
          c.id === targetChapter.id ? { ...c, doc: nextDoc, rev: c.rev + 1 } : c,