fix: seamlessly update editor collab token on expiration

Handle token refresh without disrupting editor state by disconnecting,
updating token, and reconnecting the HocuspocusProvider. This prevents
scroll position resets and maintains user context during authentication
token renewal.

apps/client/public/locales/en-US/translation.json

@@ -389,5 +389,15 @@
   "Failed to share page": "Failed to share page",
   "Copy page": "Copy page",
   "Copy page to a different space.": "Copy page to a different space.",
-  "Page copied successfully": "Page copied successfully"
+  "Page copied successfully": "Page copied successfully",
+  "Find": "Find",
+  "Not found": "Not found",
+  "Previous Match (Shift+Enter)": "Previous Match (Shift+Enter)",
+  "Next match (Enter)": "Next match (Enter)",
+  "Match case (Alt+C)": "Match case (Alt+C)",
+  "Replace": "Replace",
+  "Close (Escape)": "Close (Escape)",
+  "Replace (Enter)": "Replace (Enter)",
+  "Replace all (Ctrl+Alt+Enter)": "Replace all (Ctrl+Alt+Enter)",
+  "Replace all": "Replace all"
 }

apps/client/src/features/editor/components/search-and-replace/atoms/search-and-replace-state-atom.ts

@@ -0,0 +1,9 @@
+import { atom } from "jotai";
+
+type SearchAndReplaceAtomType = {
+  isOpen: boolean;
+};
+
+export const searchAndReplaceStateAtom = atom<SearchAndReplaceAtomType>({
+  isOpen: false,
+});

apps/client/src/features/editor/components/search-and-replace/search-and-replace-dialog.tsx

@@ -0,0 +1,312 @@
+import {
+  ActionIcon,
+  Button,
+  Dialog,
+  Flex,
+  Input,
+  Stack,
+  Text,
+  Tooltip,
+} from "@mantine/core";
+import {
+  IconArrowNarrowDown,
+  IconArrowNarrowUp,
+  IconLetterCase,
+  IconReplace,
+  IconSearch,
+  IconX,
+} from "@tabler/icons-react";
+import { useEditor } from "@tiptap/react";
+import React, { useEffect, useMemo, useRef, useState } from "react";
+import { searchAndReplaceStateAtom } from "@/features/editor/components/search-and-replace/atoms/search-and-replace-state-atom.ts";
+import { useAtom } from "jotai";
+import { useTranslation } from "react-i18next";
+import { getHotkeyHandler, useToggle } from "@mantine/hooks";
+import { useLocation } from "react-router-dom";
+import classes from "./search-replace.module.css";
+
+interface PageFindDialogDialogProps {
+  editor: ReturnType<typeof useEditor>;
+  editable?: boolean;
+}
+
+function SearchAndReplaceDialog({ editor, editable = true }: PageFindDialogDialogProps) {
+  const { t } = useTranslation();
+  const [searchText, setSearchText] = useState("");
+  const [replaceText, setReplaceText] = useState("");
+  const [pageFindState, setPageFindState] = useAtom(searchAndReplaceStateAtom);
+  const inputRef = useRef(null);
+
+  const [replaceButton, replaceButtonToggle] = useToggle([
+    { isReplaceShow: false, color: "gray" },
+    { isReplaceShow: true, color: "blue" },
+  ]);
+
+  const [caseSensitive, caseSensitiveToggle] = useToggle([
+    { isCaseSensitive: false, color: "gray" },
+    { isCaseSensitive: true, color: "blue" },
+  ]);
+
+  const searchInputEvent = (event: React.ChangeEvent<HTMLInputElement>) => {
+    setSearchText(event.target.value);
+  };
+
+  const replaceInputEvent = (event: React.ChangeEvent<HTMLInputElement>) => {
+    setReplaceText(event.target.value);
+  };
+
+  const closeDialog = () => {
+    setSearchText("");
+    setReplaceText("");
+    setPageFindState({ isOpen: false });
+    // Reset replace button state when closing
+    if (replaceButton.isReplaceShow) {
+      replaceButtonToggle();
+    }
+    // Clear search term in editor
+    if (editor) {
+      editor.commands.setSearchTerm("");
+    }
+  };
+
+  const goToSelection = () => {
+    if (!editor) return;
+
+    const { results, resultIndex } = editor.storage.searchAndReplace;
+    const position: Range = results[resultIndex];
+
+    if (!position) return;
+
+    // @ts-ignore
+    editor.commands.setTextSelection(position);
+
+    const element = document.querySelector(".search-result-current");
+    if (element)
+      element.scrollIntoView({ behavior: "smooth", block: "center" });
+
+    editor.commands.setTextSelection(0);
+  };
+
+  const next = () => {
+    editor.commands.nextSearchResult();
+    goToSelection();
+  };
+
+  const previous = () => {
+    editor.commands.previousSearchResult();
+    goToSelection();
+  };
+
+  const replace = () => {
+    editor.commands.setReplaceTerm(replaceText);
+    editor.commands.replace();
+    goToSelection();
+  };
+
+  const replaceAll = () => {
+    editor.commands.setReplaceTerm(replaceText);
+    editor.commands.replaceAll();
+  };
+
+  useEffect(() => {
+    editor.commands.setSearchTerm(searchText);
+    editor.commands.resetIndex();
+    editor.commands.selectCurrentItem();
+  }, [searchText]);
+
+  const handleOpenEvent = (e) => {
+    setPageFindState({ isOpen: true });
+    const selectedText = editor.state.doc.textBetween(
+      editor.state.selection.from,
+      editor.state.selection.to,
+    );
+    if (selectedText !== "") {
+      setSearchText(selectedText);
+    }
+    inputRef.current?.focus();
+    inputRef.current?.select();
+  };
+
+  const handleCloseEvent = (e) => {
+    closeDialog();
+  };
+
+  useEffect(() => {
+    !pageFindState.isOpen && closeDialog();
+
+    document.addEventListener("openFindDialogFromEditor", handleOpenEvent);
+    document.addEventListener("closeFindDialogFromEditor", handleCloseEvent);
+
+    return () => {
+      document.removeEventListener("openFindDialogFromEditor", handleOpenEvent);
+      document.removeEventListener(
+        "closeFindDialogFromEditor",
+        handleCloseEvent,
+      );
+    };
+  }, [pageFindState.isOpen]);
+
+  useEffect(() => {
+    editor.commands.setCaseSensitive(caseSensitive.isCaseSensitive);
+    editor.commands.resetIndex();
+    goToSelection();
+  }, [caseSensitive]);
+
+  const resultsCount = useMemo(
+    () =>
+      searchText.trim() === ""
+        ? ""
+        : editor?.storage?.searchAndReplace?.results.length > 0
+        ? editor?.storage?.searchAndReplace?.resultIndex +
+          1 +
+          "/" +
+          editor?.storage?.searchAndReplace?.results.length
+        : t("Not found"),
+    [
+      searchText,
+      editor?.storage?.searchAndReplace?.resultIndex,
+      editor?.storage?.searchAndReplace?.results.length,
+    ],
+  );
+
+  const location = useLocation();
+  useEffect(() => {
+    closeDialog();
+  }, [location]);
+
+  return (
+    <Dialog
+      className={classes.findDialog}
+      opened={pageFindState.isOpen}
+
+      size="lg"
+      radius="md"
+      w={"auto"}
+      position={{ top: 90, right: 50 }}
+      withBorder
+      transitionProps={{ transition: "slide-down" }}
+    >
+      <Stack gap="xs">
+        <Flex align="center" gap="xs">
+          <Input
+            ref={inputRef}
+            placeholder={t("Find")}
+            leftSection={<IconSearch size={16} />}
+            rightSection={
+              <Text size="xs" ta="right">
+                {resultsCount}
+              </Text>
+            }
+            rightSectionWidth="70"
+            rightSectionPointerEvents="all"
+            size="xs"
+            w={220}
+            onChange={searchInputEvent}
+            value={searchText}
+            autoFocus
+            onKeyDown={getHotkeyHandler([
+              ["Enter", next],
+              ["shift+Enter", previous],
+              ["alt+C", caseSensitiveToggle],
+              //@ts-ignore
+              ...(editable ? [["alt+R", replaceButtonToggle]] : []),
+            ])}
+          />
+
+          <ActionIcon.Group>
+            <Tooltip label={t("Previous match (Shift+Enter)")}>
+              <ActionIcon variant="subtle" color="gray" onClick={previous}>
+                <IconArrowNarrowUp
+                  style={{ width: "70%", height: "70%" }}
+                  stroke={1.5}
+                />
+              </ActionIcon>
+            </Tooltip>
+            <Tooltip label={t("Next match (Enter)")}>
+              <ActionIcon variant="subtle" color="gray" onClick={next}>
+                <IconArrowNarrowDown
+                  style={{ width: "70%", height: "70%" }}
+                  stroke={1.5}
+                />
+              </ActionIcon>
+            </Tooltip>
+            <Tooltip label={t("Match case (Alt+C)")}>
+              <ActionIcon
+                variant="subtle"
+                color={caseSensitive.color}
+                onClick={() => caseSensitiveToggle()}
+              >
+                <IconLetterCase
+                  style={{ width: "70%", height: "70%" }}
+                  stroke={1.5}
+                />
+              </ActionIcon>
+            </Tooltip>
+            {editable && (
+              <Tooltip label={t("Replace")}>
+                <ActionIcon
+                  variant="subtle"
+                  color={replaceButton.color}
+                  onClick={() => replaceButtonToggle()}
+                >
+                  <IconReplace
+                    style={{ width: "70%", height: "70%" }}
+                    stroke={1.5}
+                  />
+                </ActionIcon>
+              </Tooltip>
+            )}
+            <Tooltip label={t("Close (Escape)")}>
+              <ActionIcon variant="subtle" color="gray" onClick={closeDialog}>
+                <IconX style={{ width: "70%", height: "70%" }} stroke={1.5} />
+              </ActionIcon>
+            </Tooltip>
+          </ActionIcon.Group>
+        </Flex>
+        {replaceButton.isReplaceShow && editable && (
+          <Flex align="center" gap="xs">
+            <Input
+              placeholder={t("Replace")}
+              leftSection={<IconReplace size={16} />}
+              rightSection={<div></div>}
+              rightSectionPointerEvents="all"
+              size="xs"
+              w={180}
+              autoFocus
+              onChange={replaceInputEvent}
+              value={replaceText}
+              onKeyDown={getHotkeyHandler([
+                ["Enter", replace],
+                ["ctrl+alt+Enter", replaceAll],
+              ])}
+            />
+            <ActionIcon.Group>
+              <Tooltip label={t("Replace (Enter)")}>
+                <Button
+                  size="xs"
+                  variant="subtle"
+                  color="gray"
+                  onClick={replace}
+                >
+                  {t("Replace")}
+                </Button>
+              </Tooltip>
+              <Tooltip label={t("Replace all (Ctrl+Alt+Enter)")}>
+                <Button
+                  size="xs"
+                  variant="subtle"
+                  color="gray"
+                  onClick={replaceAll}
+                >
+                  {t("Replace all")}
+                </Button>
+              </Tooltip>
+            </ActionIcon.Group>
+          </Flex>
+        )}
+      </Stack>
+    </Dialog>
+  );
+}
+
+export default SearchAndReplaceDialog;

apps/client/src/features/editor/components/search-and-replace/search-replace.module.css

@@ -0,0 +1,10 @@
+.findDialog{
+  @media print {
+    display: none;
+  }
+}
+
+.findDialog div[data-position="right"].mantine-Input-section {
+  justify-content: right;
+  padding-right: 8px;
+}

apps/client/src/features/editor/extensions/extensions.ts

@@ -36,6 +36,7 @@ import {
   Drawio,
   Excalidraw,
   Embed,
+  SearchAndReplace,
   Mention,
 } from "@docmost/editor-ext";
 import {
@@ -217,6 +218,22 @@ export const mainExtensions = [
   CharacterCount.configure({
     wordCounter: (text) => countWords(text),
   }),
+  SearchAndReplace.extend({
+    addKeyboardShortcuts() {
+      return {
+        'Mod-f': () => {
+          const event = new CustomEvent("openFindDialogFromEditor", {});
+          document.dispatchEvent(event);
+          return true;
+        },
+        'Escape': () => {
+          const event = new CustomEvent("closeFindDialogFromEditor", {});
+          document.dispatchEvent(event);
+          return true;
+        },
+      }
+    },
+  }).configure(),
 ] as any;
 
 type CollabExtensions = (provider: HocuspocusProvider, user: IUser) => any[];

apps/client/src/features/editor/page-editor.tsx

@@ -1,10 +1,5 @@
 import "@/features/editor/styles/index.css";
-import React, {
+import React, { useEffect, useMemo, useRef, useState } from "react";
-  useEffect,
-  useMemo,
-  useRef,
-  useState,
-} from "react";
 import { IndexeddbPersistence } from "y-indexeddb";
 import * as Y from "yjs";
 import {
@@ -44,6 +39,7 @@ import LinkMenu from "@/features/editor/components/link/link-menu.tsx";
 import ExcalidrawMenu from "./components/excalidraw/excalidraw-menu";
 import DrawioMenu from "./components/drawio/drawio-menu";
 import { useCollabToken } from "@/features/auth/queries/auth-query.tsx";
+import SearchAndReplaceDialog from "@/features/editor/components/search-and-replace/search-and-replace-dialog.tsx";
 import { useDebouncedCallback, useDocumentVisibility } from "@mantine/hooks";
 import { useIdle } from "@/hooks/use-idle.ts";
 import { queryClient } from "@/main.tsx";
@@ -130,7 +126,15 @@ export default function PageEditor({
           const now = Date.now().valueOf() / 1000;
           const isTokenExpired = now >= payload.exp;
           if (isTokenExpired) {
-            refetchCollabToken();
+            refetchCollabToken().then((result) => {
+              if (result.data?.token) {
+                remote.disconnect();
+                setTimeout(() => {
+                  remote.configuration.token = result.data.token;
+                  remote.connect();
+                }, 100);
+              }
+            });
           }
         },
         onStatus: (status) => {
@@ -156,6 +160,21 @@ export default function PageEditor({
     };
   }, [pageId]);
 
+  /*
+  useEffect(() => {
+    // Handle token updates by reconnecting with new token
+    if (providersRef.current?.remote && collabQuery?.token) {
+      const currentToken = providersRef.current.remote.configuration.token;
+      if (currentToken !== collabQuery.token) {
+        // Token has changed, need to reconnect with new token
+        providersRef.current.remote.disconnect();
+        providersRef.current.remote.configuration.token = collabQuery.token;
+        providersRef.current.remote.connect();
+      }
+    }
+  }, [collabQuery?.token]);
+   */
+
   // Only connect/disconnect on tab/idle, not destroy
   useEffect(() => {
     if (!providersReady || !providersRef.current) return;
@@ -350,6 +369,11 @@ export default function PageEditor({
     <div style={{ position: "relative" }}>
       <div ref={menuContainerRef}>
         <EditorContent editor={editor} />
+
+        {editor && (
+          <SearchAndReplaceDialog editor={editor} editable={editable} />
+        )}
+
         {editor && editor.isEditable && (
           <div>
             <EditorBubbleMenu editor={editor} />

apps/client/src/features/editor/styles/details.css

@@ -71,4 +71,12 @@
   [data-type="details"][open] > [data-type="detailsButton"] .ProseMirror-icon{
     transform: rotateZ(90deg);
   }
-}
+  
+  [data-type="details"]:has(.search-result) > [data-type="detailsContainer"] > [data-type="detailsContent"]{
+    display: block;
+  }
+
+  [data-type="details"]:has(.search-result) > [data-type="detailsButton"] .ProseMirror-icon{
+    transform: rotateZ(90deg);
+  }
+}

apps/client/src/features/editor/styles/find.css

@@ -0,0 +1,9 @@
+.search-result{
+  background: #ffff65;
+  color: #212529;
+}
+
+.search-result-current{
+  background: #ffc266 !important;
+  color: #212529;
+}

apps/client/src/features/editor/styles/index.css

@@ -9,5 +9,5 @@
 @import "./media.css";
 @import "./code.css";
 @import "./print.css";
+@import "./find.css";
 @import "./mention.css";
-

apps/client/src/features/editor/title-editor.tsx

@@ -10,8 +10,11 @@ import {
   pageEditorAtom,
   titleEditorAtom,
 } from "@/features/editor/atoms/editor-atoms";
-import { updatePageData, useUpdateTitlePageMutation } from "@/features/page/queries/page-query";
+import {
-import { useDebouncedCallback } from "@mantine/hooks";
+  updatePageData,
+  useUpdateTitlePageMutation,
+} from "@/features/page/queries/page-query";
+import { useDebouncedCallback, getHotkeyHandler } from "@mantine/hooks";
 import { useAtom } from "jotai";
 import { useQueryEmit } from "@/features/websocket/use-query-emit.ts";
 import { History } from "@tiptap/extension-history";
@@ -40,7 +43,8 @@ export function TitleEditor({
   editable,
 }: TitleEditorProps) {
   const { t } = useTranslation();
-  const { mutateAsync: updateTitlePageMutationAsync } = useUpdateTitlePageMutation();
+  const { mutateAsync: updateTitlePageMutationAsync } =
+    useUpdateTitlePageMutation();
   const pageEditor = useAtomValue(pageEditorAtom);
   const [, setTitleEditor] = useAtom(titleEditorAtom);
   const emit = useQueryEmit();
@@ -108,7 +112,12 @@ export function TitleEditor({
         spaceId: page.spaceId,
         entity: ["pages"],
         id: page.id,
-        payload: { title: page.title, slugId: page.slugId, parentPageId: page.parentPageId, icon: page.icon },
+        payload: {
+          title: page.title,
+          slugId: page.slugId,
+          parentPageId: page.parentPageId,
+          icon: page.icon,
+        },
       };
 
       if (page.title !== titleEditor.getText()) return;
@@ -152,13 +161,19 @@ export function TitleEditor({
     }
   }, [userPageEditMode, titleEditor, editable]);
 
+  const openSearchDialog = () => {
+    const event = new CustomEvent("openFindDialogFromEditor", {});
+    document.dispatchEvent(event);
+  };
+
   function handleTitleKeyDown(event: any) {
     if (!titleEditor || !pageEditor || event.shiftKey) return;
-    
+
-    // Prevent focus shift when IME composition is active 
+    // Prevent focus shift when IME composition is active
     // `keyCode === 229` is added to support Safari where `isComposing` may not be reliable
-    if (event.nativeEvent.isComposing || event.nativeEvent.keyCode === 229) return;
+    if (event.nativeEvent.isComposing || event.nativeEvent.keyCode === 229)
-    
+      return;
+
     const { key } = event;
     const { $head } = titleEditor.state.selection;
 
@@ -172,5 +187,16 @@ export function TitleEditor({
     }
   }
 
-  return <EditorContent editor={titleEditor} onKeyDown={handleTitleKeyDown} />;
+  return (
+    <EditorContent
+      editor={titleEditor}
+      onKeyDown={(event) => {
+        // First handle the search hotkey
+        getHotkeyHandler([["mod+F", openSearchDialog]])(event);
+        
+        // Then handle other key events
+        handleTitleKeyDown(event);
+      }}
+    />
+  );
 }

apps/client/src/features/page/components/header/page-header-menu.tsx

@@ -9,6 +9,7 @@ import {
   IconList,
   IconMessage,
   IconPrinter,
+  IconSearch,
   IconTrash,
   IconWifiOff,
 } from "@tabler/icons-react";
@@ -16,7 +17,12 @@ import React from "react";
 import useToggleAside from "@/hooks/use-toggle-aside.tsx";
 import { useAtom } from "jotai";
 import { historyAtoms } from "@/features/page-history/atoms/history-atoms.ts";
-import { useClipboard, useDisclosure } from "@mantine/hooks";
+import {
+  getHotkeyHandler,
+  useClipboard,
+  useDisclosure,
+  useHotkeys,
+} from "@mantine/hooks";
 import { useParams } from "react-router-dom";
 import { usePageQuery } from "@/features/page/queries/page-query.ts";
 import { buildPageUrl } from "@/features/page/page.utils.ts";
@@ -32,6 +38,7 @@ import {
   pageEditorAtom,
   yjsConnectionStatusAtom,
 } from "@/features/editor/atoms/editor-atoms.ts";
+import { searchAndReplaceStateAtom } from "@/features/editor/components/search-and-replace/atoms/search-and-replace-state-atom.ts";
 import { formattedDate, timeAgo } from "@/lib/time.ts";
 import { PageStateSegmentedControl } from "@/features/user/components/page-state-pref.tsx";
 import MovePageModal from "@/features/page/components/move-page-modal.tsx";
@@ -46,6 +53,26 @@ export default function PageHeaderMenu({ readOnly }: PageHeaderMenuProps) {
   const toggleAside = useToggleAside();
   const [yjsConnectionStatus] = useAtom(yjsConnectionStatusAtom);
 
+  useHotkeys(
+    [
+      [
+        "mod+F",
+        () => {
+          const event = new CustomEvent("openFindDialogFromEditor", {});
+          document.dispatchEvent(event);
+        },
+      ],
+      [
+        "Escape",
+        () => {
+          const event = new CustomEvent("closeFindDialogFromEditor", {});
+          document.dispatchEvent(event);
+        },
+      ],
+    ],
+    [],
+  );
+
   return (
     <>
       {yjsConnectionStatus === "disconnected" && (

packages/editor-ext/src/index.ts

@@ -17,4 +17,5 @@ export * from "./lib/excalidraw";
 export * from "./lib/embed";
 export * from "./lib/mention";
 export * from "./lib/markdown";
+export * from "./lib/search-and-replace";
 export * from "./lib/embed-provider";

packages/editor-ext/src/lib/custom-code-block.ts

@@ -35,6 +35,42 @@ export const CustomCodeBlock = CodeBlockLowlight.extend<CustomCodeBlockOptions>(
             return true;
           }
         },
+        "Mod-a": () => {
+          if (this.editor.isActive("codeBlock")) {
+            const { state } = this.editor;
+            const { $from } = state.selection;
+            
+            let codeBlockNode = null;
+            let codeBlockPos = null;
+            let depth = 0;
+            
+            for (depth = $from.depth; depth > 0; depth--) {
+              const node = $from.node(depth);
+              if (node.type.name === "codeBlock") {
+                codeBlockNode = node;
+                codeBlockPos = $from.start(depth) - 1;
+                break;
+              }
+            }
+            
+            if (codeBlockNode && codeBlockPos !== null) {
+              const codeBlockStart = codeBlockPos;
+              const codeBlockEnd = codeBlockPos + codeBlockNode.nodeSize;
+              
+              const contentStart = codeBlockStart + 1;
+              const contentEnd = codeBlockEnd - 1;
+              
+              this.editor.commands.setTextSelection({
+                from: contentStart,
+                to: contentEnd,
+              });
+              
+              return true;
+            }
+          }
+          
+          return false;
+        },
       };
     },
 

packages/editor-ext/src/lib/search-and-replace/index.ts

@@ -0,0 +1,3 @@
+import { SearchAndReplace } from './search-and-replace'
+export * from './search-and-replace'
+export default SearchAndReplace

packages/editor-ext/src/lib/search-and-replace/search-and-replace.ts

@@ -0,0 +1,455 @@
+/***
+ MIT License
+ Copyright (c) 2023 - 2024 Jeet Mandaliya (Github Username: sereneinserenade)
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in all
+ copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ SOFTWARE.
+ ***/
+
+import { Extension, Range, type Dispatch } from "@tiptap/core";
+import { Decoration, DecorationSet } from "@tiptap/pm/view";
+import {
+  Plugin,
+  PluginKey,
+  type EditorState,
+  type Transaction,
+} from "@tiptap/pm/state";
+import { Node as PMNode, Mark } from "@tiptap/pm/model";
+
+declare module "@tiptap/core" {
+  interface Commands<ReturnType> {
+    search: {
+      /**
+       * @description Set search term in extension.
+       */
+      setSearchTerm: (searchTerm: string) => ReturnType;
+      /**
+       * @description Set replace term in extension.
+       */
+      setReplaceTerm: (replaceTerm: string) => ReturnType;
+      /**
+       * @description Set case sensitivity in extension.
+       */
+      setCaseSensitive: (caseSensitive: boolean) => ReturnType;
+      /**
+       * @description Reset current search result to first instance.
+       */
+      resetIndex: () => ReturnType;
+      /**
+       * @description Find next instance of search result.
+       */
+      nextSearchResult: () => ReturnType;
+      /**
+       * @description Find previous instance of search result.
+       */
+      previousSearchResult: () => ReturnType;
+      /**
+       * @description Replace first instance of search result with given replace term.
+       */
+      replace: () => ReturnType;
+      /**
+       * @description Replace all instances of search result with given replace term.
+       */
+      replaceAll: () => ReturnType;
+      /**
+       * @description Find selected instance of search result.
+       */
+      selectCurrentItem: () => ReturnType;
+    };
+  }
+}
+
+interface TextNodesWithPosition {
+  text: string;
+  pos: number;
+}
+
+const getRegex = (
+  s: string,
+  disableRegex: boolean,
+  caseSensitive: boolean,
+): RegExp => {
+  return RegExp(
+    disableRegex ? s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&") : s,
+    caseSensitive ? "gu" : "gui",
+  );
+};
+
+interface ProcessedSearches {
+  decorationsToReturn: DecorationSet;
+  results: Range[];
+}
+
+function processSearches(
+  doc: PMNode,
+  searchTerm: RegExp,
+  searchResultClass: string,
+  resultIndex: number,
+): ProcessedSearches {
+  const decorations: Decoration[] = [];
+  const results: Range[] = [];
+
+  let textNodesWithPosition: TextNodesWithPosition[] = [];
+  let index = 0;
+
+  if (!searchTerm) {
+    return {
+      decorationsToReturn: DecorationSet.empty,
+      results: [],
+    };
+  }
+
+  doc?.descendants((node, pos) => {
+    if (node.isText) {
+      if (textNodesWithPosition[index]) {
+        textNodesWithPosition[index] = {
+          text: textNodesWithPosition[index].text + node.text,
+          pos: textNodesWithPosition[index].pos,
+        };
+      } else {
+        textNodesWithPosition[index] = {
+          text: `${node.text}`,
+          pos,
+        };
+      }
+    } else {
+      index += 1;
+    }
+  });
+
+  textNodesWithPosition = textNodesWithPosition.filter(Boolean);
+
+  for (const element of textNodesWithPosition) {
+    const { text, pos } = element;
+    const matches = Array.from(text.matchAll(searchTerm)).filter(
+      ([matchText]) => matchText.trim(),
+    );
+
+    for (const m of matches) {
+      if (m[0] === "") break;
+
+      if (m.index !== undefined) {
+        results.push({
+          from: pos + m.index,
+          to: pos + m.index + m[0].length,
+        });
+      }
+    }
+  }
+
+  for (let i = 0; i < results.length; i += 1) {
+    const r = results[i];
+    const className =
+      i === resultIndex
+        ? `${searchResultClass} ${searchResultClass}-current`
+        : searchResultClass;
+    const decoration: Decoration = Decoration.inline(r.from, r.to, {
+      class: className,
+    });
+
+    decorations.push(decoration);
+  }
+
+  return {
+    decorationsToReturn: DecorationSet.create(doc, decorations),
+    results,
+  };
+}
+
+const replace = (
+  replaceTerm: string,
+  results: Range[],
+  resultIndex: number,
+  { state, dispatch }: { state: EditorState; dispatch: Dispatch },
+) => {
+  const firstResult = results[resultIndex];
+
+  if (!firstResult) return;
+
+  const { from, to } = results[resultIndex];
+
+  if (dispatch) {
+    const tr = state.tr;
+    
+    // Get all marks that span the text being replaced
+    const marksSet = new Set<Mark>();
+    state.doc.nodesBetween(from, to, (node) => {
+      if (node.isText && node.marks) {
+        node.marks.forEach(mark => marksSet.add(mark));
+      }
+    });
+    
+    const marks = Array.from(marksSet);
+    
+    // Delete the old text and insert new text with preserved marks
+    tr.delete(from, to);
+    tr.insert(from, state.schema.text(replaceTerm, marks));
+    
+    dispatch(tr);
+  }
+};
+
+const replaceAll = (
+  replaceTerm: string,
+  results: Range[],
+  { tr, dispatch }: { tr: Transaction; dispatch: Dispatch },
+) => {
+  const resultsCopy = results.slice();
+
+  if (!resultsCopy.length) return;
+
+  // Process replacements in reverse order to avoid position shifting issues
+  for (let i = resultsCopy.length - 1; i >= 0; i -= 1) {
+    const { from, to } = resultsCopy[i];
+    
+    // Get all marks that span the text being replaced
+    const marksSet = new Set<Mark>();
+    tr.doc.nodesBetween(from, to, (node) => {
+      if (node.isText && node.marks) {
+        node.marks.forEach(mark => marksSet.add(mark));
+      }
+    });
+    
+    const marks = Array.from(marksSet);
+    
+    // Delete and insert with preserved marks
+    tr.delete(from, to);
+    tr.insert(from, tr.doc.type.schema.text(replaceTerm, marks));
+  }
+
+  dispatch(tr);
+};
+
+export const searchAndReplacePluginKey = new PluginKey(
+  "searchAndReplacePlugin",
+);
+
+export interface SearchAndReplaceOptions {
+  searchResultClass: string;
+  disableRegex: boolean;
+}
+
+export interface SearchAndReplaceStorage {
+  searchTerm: string;
+  replaceTerm: string;
+  results: Range[];
+  lastSearchTerm: string;
+  caseSensitive: boolean;
+  lastCaseSensitive: boolean;
+  resultIndex: number;
+  lastResultIndex: number;
+}
+
+export const SearchAndReplace = Extension.create<
+  SearchAndReplaceOptions,
+  SearchAndReplaceStorage
+>({
+  name: "searchAndReplace",
+
+  addOptions() {
+    return {
+      searchResultClass: "search-result",
+      disableRegex: true,
+    };
+  },
+
+  addStorage() {
+    return {
+      searchTerm: "",
+      replaceTerm: "",
+      results: [],
+      lastSearchTerm: "",
+      caseSensitive: false,
+      lastCaseSensitive: false,
+      resultIndex: 0,
+      lastResultIndex: 0,
+    };
+  },
+
+  addCommands() {
+    return {
+      setSearchTerm:
+        (searchTerm: string) =>
+        ({ editor }) => {
+          editor.storage.searchAndReplace.searchTerm = searchTerm;
+
+          return false;
+        },
+      setReplaceTerm:
+        (replaceTerm: string) =>
+        ({ editor }) => {
+          editor.storage.searchAndReplace.replaceTerm = replaceTerm;
+
+          return false;
+        },
+      setCaseSensitive:
+        (caseSensitive: boolean) =>
+        ({ editor }) => {
+          editor.storage.searchAndReplace.caseSensitive = caseSensitive;
+
+          return false;
+        },
+      resetIndex:
+        () =>
+        ({ editor }) => {
+          editor.storage.searchAndReplace.resultIndex = 0;
+
+          return false;
+        },
+      nextSearchResult:
+        () =>
+        ({ editor }) => {
+          const { results, resultIndex } = editor.storage.searchAndReplace;
+
+          const nextIndex = resultIndex + 1;
+
+          if (results[nextIndex]) {
+            editor.storage.searchAndReplace.resultIndex = nextIndex;
+          } else {
+            editor.storage.searchAndReplace.resultIndex = 0;
+          }
+
+          return false;
+        },
+      previousSearchResult:
+        () =>
+        ({ editor }) => {
+          const { results, resultIndex } = editor.storage.searchAndReplace;
+
+          const prevIndex = resultIndex - 1;
+
+          if (results[prevIndex]) {
+            editor.storage.searchAndReplace.resultIndex = prevIndex;
+          } else {
+            editor.storage.searchAndReplace.resultIndex = results.length - 1;
+          }
+
+          return false;
+        },
+      replace:
+        () =>
+        ({ editor, state, dispatch }) => {
+          const { replaceTerm, results, resultIndex } =
+            editor.storage.searchAndReplace;
+
+          replace(replaceTerm, results, resultIndex, { state, dispatch });
+
+          // After replace, adjust index if needed
+          // The results will be recalculated by the plugin, but we need to ensure
+          // the index doesn't exceed the new bounds
+          setTimeout(() => {
+            const newResultsLength = editor.storage.searchAndReplace.results.length;
+            if (newResultsLength > 0 && editor.storage.searchAndReplace.resultIndex >= newResultsLength) {
+              // Keep the same position if possible, otherwise go to the last result
+              editor.storage.searchAndReplace.resultIndex = Math.min(resultIndex, newResultsLength - 1);
+            }
+          }, 0);
+
+          return false;
+        },
+      replaceAll:
+        () =>
+        ({ editor, tr, dispatch }) => {
+          const { replaceTerm, results } = editor.storage.searchAndReplace;
+
+          replaceAll(replaceTerm, results, { tr, dispatch });
+
+          return false;
+        },
+      selectCurrentItem:
+        () =>
+        ({ editor }) => {
+          const { results } = editor.storage.searchAndReplace;
+          for (let i = 0; i < results.length; i++) {
+            if (
+              results[i].from == editor.state.selection.from &&
+              results[i].to == editor.state.selection.to
+            ) {
+              editor.storage.searchAndReplace.resultIndex = i;
+            }
+          }
+          return false;
+        },
+    };
+  },
+
+  addProseMirrorPlugins() {
+    const editor = this.editor;
+    const { searchResultClass, disableRegex } = this.options;
+
+    const setLastSearchTerm = (t: string) =>
+      (editor.storage.searchAndReplace.lastSearchTerm = t);
+    const setLastCaseSensitive = (t: boolean) =>
+      (editor.storage.searchAndReplace.lastCaseSensitive = t);
+    const setLastResultIndex = (t: number) =>
+      (editor.storage.searchAndReplace.lastResultIndex = t);
+
+    return [
+      new Plugin({
+        key: searchAndReplacePluginKey,
+        state: {
+          init: () => DecorationSet.empty,
+          apply({ doc, docChanged }, oldState) {
+            const {
+              searchTerm,
+              lastSearchTerm,
+              caseSensitive,
+              lastCaseSensitive,
+              resultIndex,
+              lastResultIndex,
+            } = editor.storage.searchAndReplace;
+
+            if (
+              !docChanged &&
+              lastSearchTerm === searchTerm &&
+              lastCaseSensitive === caseSensitive &&
+              lastResultIndex === resultIndex
+            )
+              return oldState;
+
+            setLastSearchTerm(searchTerm);
+            setLastCaseSensitive(caseSensitive);
+            setLastResultIndex(resultIndex);
+
+            if (!searchTerm) {
+              editor.storage.searchAndReplace.results = [];
+              return DecorationSet.empty;
+            }
+
+            const { decorationsToReturn, results } = processSearches(
+              doc,
+              getRegex(searchTerm, disableRegex, caseSensitive),
+              searchResultClass,
+              resultIndex,
+            );
+
+            editor.storage.searchAndReplace.results = results;
+
+            return decorationsToReturn;
+          },
+        },
+        props: {
+          decorations(state) {
+            return this.getState(state);
+          },
+        },
+      }),
+    ];
+  },
+});
+
+export default SearchAndReplace;