diff --git a/schema/plugin.json b/schema/plugin.json
index 5188fca..3a5083c 100644
--- a/schema/plugin.json
+++ b/schema/plugin.json
@@ -18,22 +18,12 @@
     {
       "command": "documentsearch:end",
       "keys": ["Escape"],
-      "selector": ".jp-mod-searchable .jp-FileEditor .cm-vimMode"
+      "selector": ".jp-mod-searchable .jp-FileEditor [data-jp-vim-mode-name='normal']"
     },
     {
       "command": "documentsearch:end",
       "keys": ["Escape"],
-      "selector": ".jp-mod-searchable .jp-Notebook .cm-vimMode"
-    },
-    {
-      "command": "documentsearch:end",
-      "keys": ["Escape"],
-      "selector": ".jp-mod-search-active .jp-FileEditor .cm-vimMode"
-    },
-    {
-      "command": "documentsearch:end",
-      "keys": ["Escape"],
-      "selector": ".jp-mod-search-active .jp-Notebook .cm-vimMode"
+      "selector": ".jp-mod-search-active .jp-FileEditor [data-jp-vim-mode-name='normal']"
     },
     {
       "selector": ".jp-NotebookPanel[data-jp-vim-mode='true'] .jp-Notebook.jp-mod-editMode",
diff --git a/src/codemirrorCommands.ts b/src/codemirrorCommands.ts
index 9dc5a6c..4f2a2d0 100644
--- a/src/codemirrorCommands.ts
+++ b/src/codemirrorCommands.ts
@@ -67,13 +67,17 @@ export class VimEditorManager {
     this.modifyEditor(this._lastActiveEditor);
   }
 
+  /**
+   * Hook up vim mode into given editor.
+   * Returns true if vim mode was enabled.
+   */
   modifyEditor(editor: CodeEditor.IEditor | null): boolean {
-    // JupyterLab 4.0 only supports CodeMirror editors
-    const mirrorEditor = editor as CodeMirrorEditor | null;
-
-    if (!mirrorEditor) {
+    if (!editor) {
       throw Error('Editor not available');
     }
+    // JupyterLab 4.0 only supports CodeMirror editors
+    const mirrorEditor = editor as CodeMirrorEditor;
+
     this._lastActiveEditor = mirrorEditor;
 
     const view = mirrorEditor.editor;
@@ -89,6 +93,9 @@ export class VimEditorManager {
         // as blurred because it exists outside of the CodeMirror6 state; here
         // we override `hasFocus` handler to ensure it is taken into account.
         const cm = getCM(view)!;
+        cm.on('vim-mode-change', () => {
+          editor.host.dataset.jpVimModeName = cm.state.vim.mode;
+        });
         mirrorEditor.hasFocus = () => {
           if (
             cm.state.dialog &&