test: add tests for some components in base > prompt-editor (#32472)

Co-authored-by: sahil-infocusp <73810410+sahil-infocusp@users.noreply.github.com>

web/app/components/base/prompt-editor/hooks.spec.tsx

@@ -0,0 +1,307 @@
+import type { EntityMatch } from '@lexical/text'
+import type { Klass, LexicalEditor, TextNode } from 'lexical'
+import { render, renderHook, screen } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import { COMMAND_PRIORITY_LOW, KEY_BACKSPACE_COMMAND, KEY_DELETE_COMMAND } from 'lexical'
+import {
+  useBasicTypeaheadTriggerMatch,
+  useLexicalTextEntity,
+  useSelectOrDelete,
+  useTrigger,
+} from './hooks'
+import {
+  DELETE_CONTEXT_BLOCK_COMMAND,
+} from './plugins/context-block'
+import { ContextBlockNode } from './plugins/context-block/node'
+import { DELETE_QUERY_BLOCK_COMMAND } from './plugins/query-block'
+import { QueryBlockNode } from './plugins/query-block/node'
+
+type MockNode = {
+  isDecorator?: boolean
+  remove?: () => void
+}
+
+type MockSelection = {
+  getNodes: () => MockNode[]
+  isNodeSelection?: boolean
+}
+
+type SelectOrDeleteCommand = Parameters<typeof useSelectOrDelete>[1]
+type LexicalTextEntityGetMatch = (text: string) => null | EntityMatch
+type LexicalTextEntityCreateNode = (textNode: TextNode) => TextNode
+
+const mockState = vi.hoisted(() => {
+  const commandHandlers = new Map<unknown, (event: KeyboardEvent) => boolean>()
+  const registerCommand = vi.fn((command: unknown, handler: (event: KeyboardEvent) => boolean) => {
+    commandHandlers.set(command, handler)
+    return vi.fn()
+  })
+
+  return {
+    editor: {
+      registerCommand,
+      registerNodeTransform: vi.fn(),
+      dispatchCommand: vi.fn(),
+    },
+    commandHandlers,
+    isSelected: false,
+    setSelected: vi.fn(),
+    clearSelection: vi.fn(),
+    selection: null as MockSelection | null,
+    node: null as MockNode | null,
+    mergeRegister: vi.fn((...cleanups: Array<() => void>) => {
+      return () => {
+        cleanups.forEach(cleanup => cleanup())
+      }
+    }),
+    removePlainTextTransform: vi.fn(),
+    removeReverseNodeTransform: vi.fn(),
+  }
+})
+
+vi.mock('@lexical/react/LexicalComposerContext', () => ({
+  useLexicalComposerContext: () => [mockState.editor],
+}))
+
+vi.mock('@lexical/react/useLexicalNodeSelection', () => ({
+  useLexicalNodeSelection: () => [
+    mockState.isSelected,
+    mockState.setSelected,
+    mockState.clearSelection,
+  ],
+}))
+
+vi.mock('@lexical/utils', () => ({
+  mergeRegister: mockState.mergeRegister,
+}))
+
+vi.mock('lexical', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('lexical')>()
+  return {
+    ...actual,
+    $getSelection: () => mockState.selection,
+    $getNodeByKey: () => mockState.node,
+    $isDecoratorNode: (node: MockNode | null) => !!node?.isDecorator,
+    $isNodeSelection: (selection: MockSelection | null) => !!selection?.isNodeSelection,
+  }
+})
+
+const SelectOrDeleteHarness = ({ nodeKey, command }: {
+  nodeKey: string
+  command?: SelectOrDeleteCommand
+}) => {
+  const [ref, isSelected] = useSelectOrDelete(nodeKey, command)
+  return (
+    <div
+      ref={ref}
+      data-testid="select-or-delete-node"
+      data-selected={isSelected ? 'true' : 'false'}
+    >
+      node
+    </div>
+  )
+}
+
+const TriggerHarness = () => {
+  const [ref, open] = useTrigger()
+  return (
+    <div>
+      <div ref={ref} data-testid="trigger-target">toggle</div>
+      <span>{open ? 'open' : 'closed'}</span>
+    </div>
+  )
+}
+
+const LexicalTextEntityHarness = ({
+  getMatch,
+  targetNode,
+  createNode,
+}: {
+  getMatch: LexicalTextEntityGetMatch
+  targetNode: Klass<TextNode>
+  createNode: LexicalTextEntityCreateNode
+}) => {
+  useLexicalTextEntity(getMatch, targetNode, createNode)
+  return null
+}
+
+describe('prompt-editor/hooks', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockState.commandHandlers.clear()
+    mockState.isSelected = false
+    mockState.selection = null
+    mockState.node = null
+    mockState.editor.registerNodeTransform
+      .mockReset()
+      .mockReturnValueOnce(mockState.removePlainTextTransform)
+      .mockReturnValueOnce(mockState.removeReverseNodeTransform)
+  })
+
+  // Selection/deletion hook behavior around Lexical node commands.
+  describe('useSelectOrDelete', () => {
+    it('should register delete and backspace commands and select node on click', async () => {
+      const user = userEvent.setup()
+      render(
+        <SelectOrDeleteHarness
+          nodeKey="node-1"
+          command={DELETE_CONTEXT_BLOCK_COMMAND}
+        />,
+      )
+
+      expect(mockState.editor.registerCommand).toHaveBeenCalledWith(
+        KEY_DELETE_COMMAND,
+        expect.any(Function),
+        COMMAND_PRIORITY_LOW,
+      )
+      expect(mockState.editor.registerCommand).toHaveBeenCalledWith(
+        KEY_BACKSPACE_COMMAND,
+        expect.any(Function),
+        COMMAND_PRIORITY_LOW,
+      )
+
+      await user.click(screen.getByTestId('select-or-delete-node'))
+
+      expect(mockState.clearSelection).toHaveBeenCalled()
+      expect(mockState.setSelected).toHaveBeenCalledWith(true)
+    })
+
+    it('should dispatch delete command when unselected context block is focused', () => {
+      mockState.isSelected = false
+      mockState.selection = {
+        getNodes: () => [Object.create(ContextBlockNode.prototype) as MockNode],
+        isNodeSelection: false,
+      }
+
+      render(
+        <SelectOrDeleteHarness
+          nodeKey="node-1"
+          command={DELETE_CONTEXT_BLOCK_COMMAND}
+        />,
+      )
+
+      const deleteHandler = mockState.commandHandlers.get(KEY_DELETE_COMMAND)
+      expect(deleteHandler).toBeDefined()
+
+      const handled = deleteHandler?.(new KeyboardEvent('keydown'))
+
+      expect(handled).toBe(false)
+      expect(mockState.editor.dispatchCommand).toHaveBeenCalledWith(DELETE_CONTEXT_BLOCK_COMMAND, undefined)
+    })
+
+    it('should prevent default and remove selected decorator node on delete', () => {
+      const remove = vi.fn()
+      const preventDefault = vi.fn()
+      mockState.isSelected = true
+      mockState.selection = {
+        getNodes: () => [Object.create(QueryBlockNode.prototype) as MockNode],
+        isNodeSelection: true,
+      }
+      mockState.node = {
+        isDecorator: true,
+        remove,
+      }
+
+      render(
+        <SelectOrDeleteHarness
+          nodeKey="node-1"
+          command={DELETE_QUERY_BLOCK_COMMAND}
+        />,
+      )
+
+      const backspaceHandler = mockState.commandHandlers.get(KEY_BACKSPACE_COMMAND)
+      expect(backspaceHandler).toBeDefined()
+
+      const handled = backspaceHandler?.({ preventDefault } as unknown as KeyboardEvent)
+
+      expect(handled).toBe(true)
+      expect(preventDefault).toHaveBeenCalled()
+      expect(mockState.editor.dispatchCommand).toHaveBeenCalledWith(DELETE_QUERY_BLOCK_COMMAND, undefined)
+      expect(remove).toHaveBeenCalled()
+    })
+  })
+
+  // Trigger hook toggles dropdown/popup state from bound DOM element.
+  describe('useTrigger', () => {
+    it('should toggle open state when trigger element is clicked', async () => {
+      const user = userEvent.setup()
+      render(<TriggerHarness />)
+
+      expect(screen.getByText('closed')).toBeInTheDocument()
+
+      await user.click(screen.getByTestId('trigger-target'))
+      expect(screen.getByText('open')).toBeInTheDocument()
+
+      await user.click(screen.getByTestId('trigger-target'))
+      expect(screen.getByText('closed')).toBeInTheDocument()
+    })
+  })
+
+  // Lexical entity hook should register and cleanup transforms.
+  describe('useLexicalTextEntity', () => {
+    it('should register lexical text entity transforms and cleanup on unmount', () => {
+      class MockTargetNode {}
+      const getMatch: LexicalTextEntityGetMatch = vi.fn(() => null)
+      const createNode: LexicalTextEntityCreateNode = vi.fn((textNode: TextNode) => textNode)
+
+      const { unmount } = render(
+        <LexicalTextEntityHarness
+          getMatch={getMatch}
+          targetNode={MockTargetNode as unknown as Klass<TextNode>}
+          createNode={createNode}
+        />,
+      )
+
+      expect(mockState.editor.registerNodeTransform).toHaveBeenCalledTimes(2)
+      // Verify the first call uses TextNode, not MockTargetNode
+      const calls = mockState.editor.registerNodeTransform.mock.calls
+      expect(calls[0][0]).not.toBe(MockTargetNode)
+      expect(typeof calls[0][0]).toBe('function')
+      expect(mockState.editor.registerNodeTransform).toHaveBeenCalledWith(
+        MockTargetNode,
+        expect.any(Function),
+      )
+
+      unmount()
+
+      expect(getMatch).not.toHaveBeenCalled()
+      expect(createNode).not.toHaveBeenCalled()
+      expect(mockState.removePlainTextTransform).toHaveBeenCalled()
+      expect(mockState.removeReverseNodeTransform).toHaveBeenCalled()
+    })
+  })
+
+  // Regex trigger matcher behavior for typeahead text detection.
+  describe('useBasicTypeaheadTriggerMatch', () => {
+    it('should return match details when input satisfies trigger and length rules', () => {
+      const { result } = renderHook(() => useBasicTypeaheadTriggerMatch('@', {
+        minLength: 2,
+        maxLength: 5,
+      }))
+
+      const match = result.current('prefix @..', {} as LexicalEditor)
+      expect(match).toEqual({
+        leadOffset: 7,
+        matchingString: '..',
+        replaceableString: '@..',
+      })
+    })
+
+    it('should return null when matching text is shorter than minLength', () => {
+      const { result } = renderHook(() => useBasicTypeaheadTriggerMatch('@', {
+        minLength: 2,
+        maxLength: 5,
+      }))
+
+      expect(result.current('prefix @.', {} as LexicalEditor)).toBeNull()
+    })
+
+    it('should return null when matching text exceeds maxLength', () => {
+      const { result } = renderHook(() => useBasicTypeaheadTriggerMatch('@', {
+        minLength: 1,
+        maxLength: 2,
+      }))
+      expect(result.current('prefix @...', {} as LexicalEditor)).toBeNull()
+    })
+  })
+})

web/app/components/base/prompt-editor/index.spec.tsx

@@ -0,0 +1,269 @@
+import type { FocusEvent as ReactFocusEvent, ReactNode } from 'react'
+import type { PromptEditorProps } from './index'
+import type { ContextBlockType, HistoryBlockType } from './types'
+import { render, screen, waitFor } from '@testing-library/react'
+import { BLUR_COMMAND, FOCUS_COMMAND } from 'lexical'
+import * as React from 'react'
+import {
+  UPDATE_DATASETS_EVENT_EMITTER,
+  UPDATE_HISTORY_EVENT_EMITTER,
+} from './constants'
+import PromptEditor from './index'
+
+const mocks = vi.hoisted(() => {
+  const commandHandlers = new Map<unknown, (payload: unknown) => boolean>()
+  const subscriptions: Array<(payload: unknown) => void> = []
+  const rootElement = document.createElement('div')
+
+  return {
+    emit: vi.fn(),
+    rootLines: ['first line', 'second line'],
+    commandHandlers,
+    subscriptions,
+    rootElement,
+    editor: {
+      hasNodes: vi.fn(() => true),
+      registerCommand: vi.fn((command: unknown, handler: (payload: unknown) => boolean) => {
+        commandHandlers.set(command, handler)
+        return vi.fn()
+      }),
+      registerUpdateListener: vi.fn(() => vi.fn()),
+      dispatchCommand: vi.fn(),
+      getRootElement: vi.fn(() => rootElement),
+      parseEditorState: vi.fn(() => ({ state: 'parsed' })),
+      setEditorState: vi.fn(),
+      focus: vi.fn(),
+      update: vi.fn((fn: () => void) => fn()),
+    },
+  }
+})
+
+vi.mock('@/context/event-emitter', () => ({
+  useEventEmitterContextContext: () => ({
+    eventEmitter: {
+      emit: mocks.emit,
+      useSubscription: (cb: (payload: unknown) => void) => {
+        mocks.subscriptions.push(cb)
+      },
+    },
+  }),
+}))
+
+vi.mock('@lexical/code', () => ({
+  CodeNode: class CodeNode {},
+}))
+
+vi.mock('@lexical/react/LexicalComposerContext', () => ({
+  useLexicalComposerContext: () => [mocks.editor],
+}))
+
+vi.mock('lexical', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('lexical')>()
+  return {
+    ...actual,
+    $getRoot: () => ({
+      getChildren: () => mocks.rootLines.map(line => ({
+        getTextContent: () => line,
+      })),
+    }),
+    TextNode: class TextNode {
+      __text: string
+      constructor(text = '') {
+        this.__text = text
+      }
+    },
+  }
+})
+
+vi.mock('@lexical/react/LexicalComposer', () => ({
+  LexicalComposer: ({ children }: { children: ReactNode }) => (
+    <div data-testid="lexical-composer">{children}</div>
+  ),
+}))
+
+vi.mock('@lexical/react/LexicalContentEditable', () => ({
+  ContentEditable: (props: React.HTMLAttributes<HTMLDivElement>) => <div data-testid="content-editable" {...props} />,
+}))
+
+vi.mock('@lexical/react/LexicalErrorBoundary', () => ({
+  LexicalErrorBoundary: () => <div data-testid="lexical-error-boundary" />,
+}))
+
+vi.mock('@lexical/react/LexicalHistoryPlugin', () => ({
+  HistoryPlugin: () => <div data-testid="history-plugin" />,
+}))
+
+vi.mock('@lexical/react/LexicalOnChangePlugin', () => ({
+  OnChangePlugin: ({ onChange }: { onChange: (editorState: { read: (fn: () => void) => void }) => void }) => {
+    React.useEffect(() => {
+      onChange({
+        read: (fn: () => void) => fn(),
+      })
+    }, [onChange])
+    return <div data-testid="on-change-plugin" />
+  },
+}))
+
+vi.mock('@lexical/react/LexicalRichTextPlugin', () => ({
+  RichTextPlugin: ({ contentEditable, placeholder }: { contentEditable: ReactNode, placeholder: ReactNode }) => (
+    <div data-testid="rich-text-plugin">
+      {contentEditable}
+      {placeholder}
+    </div>
+  ),
+}))
+
+vi.mock('@lexical/react/LexicalTypeaheadMenuPlugin', () => ({
+  MenuOption: class MenuOption {
+    key: string
+    constructor(key: string) {
+      this.key = key
+    }
+  },
+  LexicalTypeaheadMenuPlugin: () => <div data-testid="typeahead-plugin" />,
+}))
+
+vi.mock('@lexical/react/LexicalDraggableBlockPlugin', () => ({
+  DraggableBlockPlugin_EXPERIMENTAL: ({ menuComponent, targetLineComponent }: {
+    menuComponent: ReactNode
+    targetLineComponent: ReactNode
+  }) => (
+    <div data-testid="draggable-plugin">
+      {menuComponent}
+      {targetLineComponent}
+    </div>
+  ),
+}))
+
+describe('PromptEditor', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mocks.commandHandlers.clear()
+    mocks.subscriptions.length = 0
+    mocks.rootLines = ['first line', 'second line']
+  })
+
+  // Rendering shell and text output from lexical state.
+  describe('Rendering', () => {
+    it('should render placeholder and call onChange with joined lexical text', async () => {
+      const onChange = vi.fn()
+
+      render(
+        <PromptEditor
+          compact={true}
+          className="editor-class"
+          placeholder="Type prompt"
+          value="seed-value"
+          onChange={onChange}
+        />,
+      )
+
+      expect(screen.getByText('Type prompt')).toBeInTheDocument()
+      expect(screen.getByTestId('content-editable')).toHaveClass('editor-class')
+      expect(screen.getByTestId('content-editable')).toHaveClass('text-[13px]')
+
+      await waitFor(() => {
+        expect(onChange).toHaveBeenCalledWith('first line\nsecond line')
+      })
+    })
+  })
+
+  // Event emitter integration for datasets and history updates.
+  describe('Event Emission', () => {
+    it('should emit dataset and history updates when corresponding props change', () => {
+      const contextBlock: ContextBlockType = {
+        show: false,
+        datasets: [{ id: 'ds-1', name: 'Dataset One', type: 'dataset' }],
+      }
+      const historyBlock: HistoryBlockType = {
+        show: false,
+        history: { user: 'user-role', assistant: 'assistant-role' },
+      }
+
+      const { rerender } = render(
+        <PromptEditor
+          contextBlock={contextBlock}
+          historyBlock={historyBlock}
+        />,
+      )
+
+      expect(mocks.emit).toHaveBeenCalledWith({
+        type: UPDATE_DATASETS_EVENT_EMITTER,
+        payload: [{ id: 'ds-1', name: 'Dataset One', type: 'dataset' }],
+      })
+      expect(mocks.emit).toHaveBeenCalledWith({
+        type: UPDATE_HISTORY_EVENT_EMITTER,
+        payload: { user: 'user-role', assistant: 'assistant-role' },
+      })
+
+      rerender(
+        <PromptEditor
+          contextBlock={{
+            show: false,
+            datasets: [{ id: 'ds-2', name: 'Dataset Two', type: 'dataset' }],
+          }}
+          historyBlock={{
+            show: false,
+            history: { user: 'user-next', assistant: 'assistant-next' },
+          }}
+        />,
+      )
+
+      expect(mocks.emit).toHaveBeenCalledWith({
+        type: UPDATE_DATASETS_EVENT_EMITTER,
+        payload: [{ id: 'ds-2', name: 'Dataset Two', type: 'dataset' }],
+      })
+      expect(mocks.emit).toHaveBeenCalledWith({
+        type: UPDATE_HISTORY_EVENT_EMITTER,
+        payload: { user: 'user-next', assistant: 'assistant-next' },
+      })
+    })
+  })
+
+  // OnBlurBlock command callbacks should forward to PromptEditor handlers.
+  describe('Focus/Blur Callbacks', () => {
+    it('should call onFocus and onBlur when lexical focus/blur commands fire', () => {
+      const onFocus = vi.fn()
+      const onBlur = vi.fn()
+
+      render(
+        <PromptEditor
+          onFocus={onFocus}
+          onBlur={onBlur}
+        />,
+      )
+
+      const focusHandler = mocks.commandHandlers.get(FOCUS_COMMAND)
+      const blurHandler = mocks.commandHandlers.get(BLUR_COMMAND)
+
+      expect(focusHandler).toBeDefined()
+      expect(blurHandler).toBeDefined()
+
+      focusHandler?.(undefined)
+      blurHandler?.({ relatedTarget: null } as ReactFocusEvent<Element>)
+
+      expect(onFocus).toHaveBeenCalledTimes(1)
+      expect(onBlur).toHaveBeenCalledTimes(1)
+    })
+  })
+
+  // Prop typing guard for shortcut popup shape without any-casts.
+  describe('Props Typing', () => {
+    it('should accept typed shortcut popup configuration', () => {
+      const Popup: NonNullable<PromptEditorProps['shortcutPopups']>[number]['Popup'] = ({ onClose }) => (
+        <button type="button" onClick={onClose}>close</button>
+      )
+
+      render(
+        <PromptEditor
+          shortcutPopups={[{
+            hotkey: ['mod', '/'],
+            Popup,
+          }]}
+        />,
+      )
+
+      expect(screen.getByTestId('lexical-composer')).toBeInTheDocument()
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/__tests__/utils.ts

@@ -0,0 +1,19 @@
+import type { Klass, LexicalEditor, LexicalNode } from 'lexical'
+import { createEditor } from 'lexical'
+
+export function createTestEditor(nodes: Array<Klass<LexicalNode>> = []) {
+  const editor = createEditor({
+    nodes,
+    onError: (error) => { throw error },
+  })
+  const root = document.createElement('div')
+  editor.setRootElement(root)
+  return editor
+}
+
+export function withEditorUpdate(
+  editor: LexicalEditor,
+  fn: () => void,
+) {
+  editor.update(fn, { discrete: true })
+}

web/app/components/base/prompt-editor/plugins/context-block/component.spec.tsx

@@ -0,0 +1,398 @@
+import { act, render, screen } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import { UPDATE_DATASETS_EVENT_EMITTER } from '../../constants'
+import ContextBlockComponent from './component'
+// Mock the hooks used by ContextBlockComponent
+const mockUseSelectOrDelete = vi.fn()
+const mockUseTrigger = vi.fn()
+
+vi.mock('../../hooks', () => ({
+  useSelectOrDelete: (...args: unknown[]) => mockUseSelectOrDelete(...args),
+  useTrigger: (...args: unknown[]) => mockUseTrigger(...args),
+}))
+
+// Mock event emitter context
+const mockUseSubscription = vi.fn()
+vi.mock('@/context/event-emitter', () => ({
+  useEventEmitterContextContext: () => ({
+    eventEmitter: {
+      useSubscription: mockUseSubscription,
+    },
+  }),
+}))
+
+// Helpers
+const defaultSetup = (overrides?: { isSelected?: boolean, open?: boolean }) => {
+  const triggerSetOpen = vi.fn()
+  mockUseSelectOrDelete.mockReturnValue([{ current: null }, overrides?.isSelected ?? false])
+  mockUseTrigger.mockReturnValue([{ current: null }, overrides?.open ?? false, triggerSetOpen])
+  return { triggerSetOpen }
+}
+
+const mockDatasets = [
+  { id: '1', name: 'Dataset A', type: 'text' },
+  { id: '2', name: 'Dataset B', type: 'text' },
+]
+
+describe('ContextBlockComponent', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  describe('Rendering', () => {
+    it('should render without crashing', () => {
+      defaultSetup()
+      const { container } = render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(container.firstChild).toBeInTheDocument()
+    })
+
+    it('should display the context title', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(screen.getByText('common.promptEditor.context.item.title')).toBeInTheDocument()
+    })
+
+    it('should display the dataset count', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(screen.getByText('2')).toBeInTheDocument()
+    })
+
+    it('should display zero count when no datasets provided', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(screen.getByText('0')).toBeInTheDocument()
+    })
+
+    it('should render the file icon', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      // File05 icon renders as an SVG
+      const fileIcon = screen.getByTestId('file-icon')
+      expect(fileIcon).toBeInTheDocument()
+    })
+  })
+
+  describe('Props', () => {
+    it('should apply selected border class when isSelected is true', () => {
+      defaultSetup({ isSelected: true })
+      const { container } = render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(container.firstChild).toHaveClass('!border-[#9B8AFB]')
+    })
+
+    it('should not apply selected border class when isSelected is false', () => {
+      defaultSetup({ isSelected: false })
+      const { container } = render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(container.firstChild).not.toHaveClass('!border-[#9B8AFB]')
+    })
+
+    it('should apply open background class when dropdown is open', () => {
+      defaultSetup({ open: true })
+      const { container } = render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(container.firstChild).toHaveClass('bg-[#EBE9FE]')
+    })
+
+    it('should apply default background class when dropdown is closed', () => {
+      defaultSetup({ open: false })
+      const { container } = render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(container.firstChild).toHaveClass('bg-[#F4F3FF]')
+    })
+
+    it('should hide the portal trigger when canNotAddContext is true', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+          canNotAddContext
+        />,
+      )
+      // The dataset count badge should not be rendered
+      expect(screen.queryByText('2')).not.toBeInTheDocument()
+    })
+  })
+
+  describe('Dropdown Content', () => {
+    it('should show dataset list when dropdown is open', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(screen.getByText('Dataset A')).toBeInTheDocument()
+      expect(screen.getByText('Dataset B')).toBeInTheDocument()
+    })
+
+    it('should show modal title with dataset count when open', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(
+        screen.getByText(/common\.promptEditor\.context\.modal\.title/),
+      ).toBeInTheDocument()
+    })
+
+    it('should show the add context button when open', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(
+        screen.getByText('common.promptEditor.context.modal.add'),
+      ).toBeInTheDocument()
+    })
+
+    it('should show the footer text when open', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(
+        screen.getByText('common.promptEditor.context.modal.footer'),
+      ).toBeInTheDocument()
+    })
+
+    it('should render folder icon for each dataset', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      const folders = screen.getAllByTestId('folder-icon')
+      expect(folders.length).toBeGreaterThanOrEqual(2)
+    })
+
+    it('should not render dropdown content when canNotAddContext is true', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+          canNotAddContext
+        />,
+      )
+      // Modal content should not be present
+      expect(screen.queryByText('Dataset A')).not.toBeInTheDocument()
+      expect(
+        screen.queryByText('common.promptEditor.context.modal.add'),
+      ).not.toBeInTheDocument()
+    })
+  })
+
+  describe('User Interactions', () => {
+    it('should call onAddContext when add button is clicked', async () => {
+      defaultSetup({ open: true })
+      const handleAddContext = vi.fn()
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={handleAddContext}
+        />,
+      )
+
+      const addButton = screen.getByTestId('add-button')
+      await userEvent.click(addButton)
+      expect(handleAddContext).toHaveBeenCalledTimes(1)
+    })
+
+    it('should render the count badge with open styles when dropdown is open', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      const countBadge = screen.getByText('2')
+      expect(countBadge).toHaveClass('bg-[#6938EF]')
+      expect(countBadge).toHaveClass('text-white')
+    })
+
+    it('should render the count badge with closed styles when dropdown is closed', () => {
+      defaultSetup({ open: false })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+      const countBadge = screen.getByText('2')
+      expect(countBadge).toHaveClass('bg-white/50')
+    })
+  })
+
+  describe('Event Emitter Subscription', () => {
+    it('should subscribe to event emitter on mount', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(mockUseSubscription).toHaveBeenCalled()
+    })
+
+    it('should update local datasets when UPDATE_DATASETS_EVENT_EMITTER event fires', () => {
+      defaultSetup({ open: true })
+      // Capture the subscription callback
+      let subscriptionCallback: (v: Record<string, unknown>) => void = () => { }
+      mockUseSubscription.mockImplementation((cb: (v: Record<string, unknown>) => void) => {
+        subscriptionCallback = cb
+      })
+
+      const { rerender } = render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={[]}
+          onAddContext={vi.fn()}
+        />,
+      )
+
+      // Initially no datasets
+      expect(screen.getByText('0')).toBeInTheDocument()
+
+      // Simulate event with new datasets
+      act(() => {
+        subscriptionCallback({
+          type: UPDATE_DATASETS_EVENT_EMITTER,
+          payload: [
+            { id: '3', name: 'New Dataset', type: 'text' },
+          ],
+        })
+      })
+
+      // Re-render to see state updates
+      rerender(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={[]}
+          onAddContext={vi.fn()}
+        />,
+      )
+
+      expect(screen.getByText('1')).toBeInTheDocument()
+      expect(screen.getByText('New Dataset')).toBeInTheDocument()
+    })
+
+    it('should not update datasets when event type does not match', () => {
+      defaultSetup({ open: true })
+      let subscriptionCallback: (v: Record<string, unknown>) => void = () => { }
+      mockUseSubscription.mockImplementation((cb: (v: Record<string, unknown>) => void) => {
+        subscriptionCallback = cb
+      })
+
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={mockDatasets}
+          onAddContext={vi.fn()}
+        />,
+      )
+
+      // Fire a different event
+      act(() => {
+        subscriptionCallback({
+          type: 'some-other-event',
+          payload: [{ id: '3', name: 'Should Not Appear', type: 'text' }],
+        })
+      })
+
+      expect(screen.queryByText('Should Not Appear')).not.toBeInTheDocument()
+      // Original datasets still there
+      expect(screen.getByText('Dataset A')).toBeInTheDocument()
+    })
+  })
+
+  describe('Edge Cases', () => {
+    it('should handle empty datasets array', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={[]}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(screen.getByText('0')).toBeInTheDocument()
+    })
+
+    it('should default datasets to empty array when undefined', () => {
+      defaultSetup()
+      render(
+        <ContextBlockComponent nodeKey="test-key" onAddContext={vi.fn()} />,
+      )
+      expect(screen.getByText('0')).toBeInTheDocument()
+    })
+
+    it('should handle single dataset', () => {
+      defaultSetup({ open: true })
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={[{ id: '1', name: 'Single', type: 'text' }]}
+          onAddContext={vi.fn()}
+        />,
+      )
+      expect(screen.getByText('1')).toBeInTheDocument()
+      expect(screen.getByText('Single')).toBeInTheDocument()
+    })
+
+    it('should handle dataset with long name by truncating', () => {
+      defaultSetup({ open: true })
+      const longName = 'A'.repeat(200)
+      render(
+        <ContextBlockComponent
+          nodeKey="test-key"
+          datasets={[{ id: '1', name: longName, type: 'text' }]}
+          onAddContext={vi.fn()}
+        />,
+      )
+      const nameElement = screen.getByText(longName)
+      expect(nameElement).toHaveClass('truncate')
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/context-block/component.tsx

@@ -1,11 +1,8 @@
 import type { FC } from 'react'
 import type { Dataset } from './index'
-import {
-  RiAddLine,
-} from '@remixicon/react'
+
 import { useState } from 'react'
 import { useTranslation } from 'react-i18next'
-import { File05, Folder } from '@/app/components/base/icons/src/vender/solid/files'
 import {
   PortalToFollowElem,
   PortalToFollowElemContent,
@@ -44,12 +41,12 @@ const ContextBlockComponent: FC<ContextBlockComponentProps> = ({
     <div
       className={`
       group inline-flex h-6 items-center rounded-[5px] border border-transparent bg-[#F4F3FF] pl-1 pr-0.5 text-[#6938EF] hover:bg-[#EBE9FE]
-      ${open ? 'bg-[#EBE9FE]' : 'bg-[#F4F3FF]'}
+      ${open ? 'bg-[#EBE9FE]' : ''}
       ${isSelected && '!border-[#9B8AFB]'}
     `}
       ref={ref}
     >
-      <File05 className="mr-1 h-[14px] w-[14px]" />
+      <span className="i-custom-vender-solid-files-file-05 mr-1 h-[14px] w-[14px]" data-testid="file-icon" />
       <div className="mr-1 text-xs font-medium">{t('promptEditor.context.item.title', { ns: 'common' })}</div>
       {!canNotAddContext && (
         <PortalToFollowElem
@@ -80,7 +77,7 @@ const ContextBlockComponent: FC<ContextBlockComponentProps> = ({
                     localDatasets.map(dataset => (
                       <div key={dataset.id} className="flex h-8 items-center">
                         <div className="mr-2 flex h-6 w-6 shrink-0 items-center justify-center rounded-md border-[0.5px] border-[#EAECF5] bg-[#F5F8FF]">
-                          <Folder className="h-4 w-4 text-[#444CE7]" />
+                          <span className="i-custom-vender-solid-files-folder h-4 w-4 text-[#444CE7]" data-testid="folder-icon" />
                         </div>
                         <div className="truncate text-sm text-gray-800" title="">{dataset.name}</div>
                       </div>
@@ -88,8 +85,8 @@ const ContextBlockComponent: FC<ContextBlockComponentProps> = ({
                   }
                 </div>
                 <div className="flex h-8 cursor-pointer items-center text-[#155EEF]" onClick={onAddContext}>
-                  <div className="mr-2 flex h-6 w-6 shrink-0 items-center justify-center rounded-md border-[0.5px] border-gray-100">
-                    <RiAddLine className="h-[14px] w-[14px]" />
+                  <div className="mr-2 flex h-6 w-6 shrink-0 items-center justify-center rounded-md border-[0.5px] border-gray-100" data-testid="add-button">
+                    <span className="i-ri-add-line h-[14px] w-[14px]" />
                   </div>
                   <div className="text-[13px] font-medium" title="">{t('promptEditor.context.modal.add', { ns: 'common' })}</div>
                 </div>

web/app/components/base/prompt-editor/plugins/context-block/context-block-replacement-block.spec.tsx

@@ -0,0 +1,296 @@
+import type { LexicalEditor } from 'lexical'
+import type { ReactNode } from 'react'
+import { LexicalComposer } from '@lexical/react/LexicalComposer'
+import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
+import { render } from '@testing-library/react'
+import { $createParagraphNode, $getRoot, $nodesOfType } from 'lexical'
+import * as React from 'react'
+import { ContextBlockNode } from '../context-block/node'
+import { $createCustomTextNode, CustomTextNode } from '../custom-text/node'
+import ContextBlockReplacementBlock from './context-block-replacement-block'
+
+// Mock the component rendered by ContextBlockNode.decorate()
+vi.mock('./component', () => ({
+  default: () => null,
+}))
+
+function createEditorConfig() {
+  return {
+    namespace: 'test',
+    nodes: [CustomTextNode, ContextBlockNode],
+    onError: (error: Error) => { throw error },
+  }
+}
+
+function TestWrapper({ children }: { children: ReactNode }) {
+  return (
+    <LexicalComposer initialConfig={createEditorConfig()}>
+      {children}
+    </LexicalComposer>
+  )
+}
+
+function renderWithEditor(ui: ReactNode) {
+  return render(ui, { wrapper: TestWrapper })
+}
+
+// Captures the editor instance so we can do updates after the initial render
+let capturedEditor: LexicalEditor | null = null
+
+const defaultOnCapture = (editor: LexicalEditor) => {
+  capturedEditor = editor
+}
+
+function EditorCapture({ onCapture = defaultOnCapture }: { onCapture?: (e: LexicalEditor) => void }) {
+  const [editor] = useLexicalComposerContext()
+  React.useEffect(() => {
+    onCapture(editor)
+  }, [editor, onCapture])
+  return null
+}
+
+type ReadResult = {
+  count: number
+  datasets: Array<{ id: string, name: string, type: string }>
+  canNotAddContext: boolean
+}
+
+function insertTextAndRead(text: string): ReadResult {
+  if (!capturedEditor)
+    throw new Error('Editor not captured')
+
+  // Insert CustomTextNode with the given text
+  capturedEditor.update(() => {
+    const root = $getRoot()
+    root.clear()
+    const paragraph = $createParagraphNode()
+    const textNode = $createCustomTextNode(text)
+    paragraph.append(textNode)
+    root.append(paragraph)
+  }, { discrete: true })
+
+  // Read the resulting state — extract all properties inside .read()
+  const result: ReadResult = { count: 0, datasets: [], canNotAddContext: false }
+  capturedEditor.getEditorState().read(() => {
+    const nodes = $nodesOfType(ContextBlockNode)
+    result.count = nodes.length
+    if (nodes.length > 0) {
+      result.datasets = nodes[0].getDatasets()
+      result.canNotAddContext = nodes[0].getCanNotAddContext()
+    }
+  })
+  return result
+}
+
+describe('ContextBlockReplacementBlock', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    capturedEditor = null
+  })
+
+  describe('Rendering', () => {
+    it('should render without crashing', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+      expect(capturedEditor).not.toBeNull()
+    })
+
+    it('should return null (no visible output from the plugin itself)', () => {
+      const { container } = renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+      expect(container.querySelector('[data-testid]')).toBeNull()
+    })
+  })
+
+  describe('Editor Node Registration Check', () => {
+    it('should not throw when ContextBlockNode is registered', () => {
+      expect(() => {
+        renderWithEditor(
+          <>
+            <ContextBlockReplacementBlock />
+            <EditorCapture />
+          </>,
+        )
+      }).not.toThrow()
+    })
+
+    it('should throw when ContextBlockNode is not registered', () => {
+      const configWithoutNode = {
+        namespace: 'test',
+        nodes: [CustomTextNode],
+        onError: (error: Error) => { throw error },
+      }
+
+      expect(() => {
+        render(
+          <LexicalComposer initialConfig={configWithoutNode}>
+            <ContextBlockReplacementBlock />
+          </LexicalComposer>,
+        )
+      }).toThrow('ContextBlockNodePlugin: ContextBlockNode not registered on editor')
+    })
+  })
+
+  describe('Text Replacement Transform', () => {
+    it('should replace context placeholder text with a ContextBlockNode', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('{{#context#}}')
+      expect(result.count).toBe(1)
+    })
+
+    it('should not replace text that is not the placeholder', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('just some normal text')
+      expect(result.count).toBe(0)
+    })
+
+    it('should not replace partial placeholder text', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('{{#contex')
+      expect(result.count).toBe(0)
+    })
+
+    it('should pass datasets to the created ContextBlockNode', () => {
+      const datasets = [{ id: '1', name: 'Test', type: 'text' }]
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock datasets={datasets} onAddContext={vi.fn()} />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('{{#context#}}')
+      expect(result.count).toBe(1)
+      expect(result.datasets).toEqual(datasets)
+    })
+
+    it('should pass canNotAddContext to the created ContextBlockNode', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock canNotAddContext={true} />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('{{#context#}}')
+      expect(result.count).toBe(1)
+      expect(result.canNotAddContext).toBe(true)
+    })
+  })
+
+  describe('onInsert callback', () => {
+    it('should call onInsert when a placeholder is replaced', () => {
+      const onInsert = vi.fn()
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock onInsert={onInsert} />
+          <EditorCapture />
+        </>,
+      )
+
+      insertTextAndRead('{{#context#}}')
+      expect(onInsert).toHaveBeenCalledTimes(1)
+    })
+
+    it('should not call onInsert when no placeholder is found', () => {
+      const onInsert = vi.fn()
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock onInsert={onInsert} />
+          <EditorCapture />
+        </>,
+      )
+
+      insertTextAndRead('no placeholder here')
+      expect(onInsert).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('Props Defaults', () => {
+    it('should default datasets to empty array', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('{{#context#}}')
+      expect(result.datasets).toEqual([])
+    })
+
+    it('should default canNotAddContext to false', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('{{#context#}}')
+      expect(result.canNotAddContext).toBe(false)
+    })
+  })
+
+  describe('Edge Cases', () => {
+    it('should handle undefined datasets prop', () => {
+      expect(() => {
+        renderWithEditor(
+          <>
+            <ContextBlockReplacementBlock datasets={undefined} />
+            <EditorCapture />
+          </>,
+        )
+      }).not.toThrow()
+    })
+
+    it('should handle empty datasets array', () => {
+      expect(() => {
+        renderWithEditor(
+          <>
+            <ContextBlockReplacementBlock datasets={[]} />
+            <EditorCapture />
+          </>,
+        )
+      }).not.toThrow()
+    })
+
+    it('should handle empty string text', () => {
+      renderWithEditor(
+        <>
+          <ContextBlockReplacementBlock />
+          <EditorCapture />
+        </>,
+      )
+
+      const result = insertTextAndRead('')
+      expect(result.count).toBe(0)
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/context-block/index.spec.tsx

@@ -0,0 +1,236 @@
+import type { LexicalEditor } from 'lexical'
+import type { ReactNode } from 'react'
+import type { Dataset } from './index'
+import { LexicalComposer } from '@lexical/react/LexicalComposer'
+import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
+import { render } from '@testing-library/react'
+import { $createParagraphNode, $getRoot } from 'lexical'
+import * as React from 'react'
+import { ContextBlock, DELETE_CONTEXT_BLOCK_COMMAND, INSERT_CONTEXT_BLOCK_COMMAND } from './index'
+import { ContextBlockNode } from './node'
+
+const mockCreateContextBlockNode = vi.fn()
+
+vi.mock('./node', async () => {
+  const actual = await vi.importActual<typeof import('./node')>('./node')
+
+  return {
+    ...actual,
+    $createContextBlockNode: (datasets: Dataset[], onAddContext: () => void, canNotAddContext?: boolean) => {
+      mockCreateContextBlockNode(datasets, onAddContext, canNotAddContext)
+      return actual.$createContextBlockNode(datasets, onAddContext, canNotAddContext)
+    },
+  }
+})
+
+vi.mock('./component', () => ({
+  default: () => null,
+}))
+
+type EditorConfig = {
+  namespace: string
+  nodes: [typeof ContextBlockNode] | []
+  onError: (error: Error) => void
+}
+
+function createEditorConfig(includeContextBlockNode = true): EditorConfig {
+  return {
+    namespace: 'test',
+    nodes: includeContextBlockNode ? [ContextBlockNode] : [],
+    onError: (error: Error) => { throw error },
+  }
+}
+
+let capturedEditor: LexicalEditor | null = null
+
+function EditorCapture() {
+  const [editor] = useLexicalComposerContext()
+  React.useEffect(() => {
+    capturedEditor = editor
+  }, [editor])
+  return null
+}
+
+function renderWithEditor(ui: ReactNode, includeContextBlockNode = true) {
+  return render(
+    <LexicalComposer initialConfig={createEditorConfig(includeContextBlockNode)}>
+      {ui}
+      <EditorCapture />
+    </LexicalComposer>,
+  )
+}
+
+function setupParagraphSelection() {
+  if (!capturedEditor)
+    throw new Error('Editor not captured')
+
+  capturedEditor.update(() => {
+    const root = $getRoot()
+    root.clear()
+    const paragraph = $createParagraphNode()
+    root.append(paragraph)
+    paragraph.select()
+  }, { discrete: true })
+}
+
+function dispatchInsert() {
+  if (!capturedEditor)
+    throw new Error('Editor not captured')
+
+  setupParagraphSelection()
+  return capturedEditor.dispatchCommand(INSERT_CONTEXT_BLOCK_COMMAND, undefined)
+}
+
+function dispatchDelete() {
+  if (!capturedEditor)
+    throw new Error('Editor not captured')
+
+  return capturedEditor.dispatchCommand(DELETE_CONTEXT_BLOCK_COMMAND, undefined)
+}
+
+describe('ContextBlock', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    capturedEditor = null
+  })
+
+  describe('Rendering', () => {
+    it('should render (no visible output)', () => {
+      const { container } = renderWithEditor(<ContextBlock />)
+      expect(container.childElementCount).toBe(0)
+    })
+  })
+
+  describe('Editor Node Registration Check', () => {
+    it('should not throw when ContextBlockNode is registered', () => {
+      expect(() => {
+        renderWithEditor(<ContextBlock />)
+      }).not.toThrow()
+    })
+
+    it('should throw when ContextBlockNode is not registered', () => {
+      expect(() => {
+        renderWithEditor(<ContextBlock />, false)
+      }).toThrow('ContextBlockPlugin: ContextBlock not registered on editor')
+    })
+  })
+
+  describe('INSERT_CONTEXT_BLOCK_COMMAND handler', () => {
+    it('should insert a context block node with default props', () => {
+      renderWithEditor(<ContextBlock />)
+
+      const handled = dispatchInsert()
+
+      expect(handled).toBe(true)
+      expect(mockCreateContextBlockNode).toHaveBeenCalledWith([], expect.any(Function), undefined)
+    })
+
+    it('should call onInsert when provided', () => {
+      const onInsert = vi.fn()
+      renderWithEditor(<ContextBlock onInsert={onInsert} />)
+
+      dispatchInsert()
+
+      expect(onInsert).toHaveBeenCalledTimes(1)
+    })
+
+    it('should pass datasets to the created node', () => {
+      const datasets: Dataset[] = [{ id: '1', name: 'Test', type: 'text' }]
+      renderWithEditor(<ContextBlock datasets={datasets} />)
+
+      dispatchInsert()
+      expect(mockCreateContextBlockNode).toHaveBeenCalledWith(datasets, expect.any(Function), undefined)
+    })
+
+    it('should pass canNotAddContext to the created node', () => {
+      renderWithEditor(<ContextBlock canNotAddContext={true} />)
+
+      dispatchInsert()
+      expect(mockCreateContextBlockNode).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.any(Function),
+        true,
+      )
+    })
+  })
+
+  describe('DELETE_CONTEXT_BLOCK_COMMAND handler', () => {
+    it('should return true when dispatched', () => {
+      renderWithEditor(<ContextBlock />)
+
+      const handled = dispatchDelete()
+
+      expect(handled).toBe(true)
+    })
+
+    it('should call onDelete when provided', () => {
+      const onDelete = vi.fn()
+      renderWithEditor(<ContextBlock onDelete={onDelete} />)
+
+      dispatchDelete()
+
+      expect(onDelete).toHaveBeenCalledTimes(1)
+    })
+
+    it('should not throw when onDelete is not provided', () => {
+      renderWithEditor(<ContextBlock />)
+
+      expect(() => dispatchDelete()).not.toThrow()
+    })
+  })
+
+  describe('Props Defaults', () => {
+    it('should default onAddContext to a noop function', () => {
+      renderWithEditor(<ContextBlock />)
+
+      dispatchInsert()
+      const onAddContextArg = mockCreateContextBlockNode.mock.calls[0][1] as () => void
+
+      expect(typeof onAddContextArg).toBe('function')
+      expect(() => onAddContextArg()).not.toThrow()
+    })
+  })
+
+  describe('Lifecycle', () => {
+    it('should unregister commands on unmount', () => {
+      const onDelete = vi.fn()
+      const { unmount } = renderWithEditor(<ContextBlock onDelete={onDelete} />)
+
+      unmount()
+      const handledAfterUnmount = dispatchDelete()
+
+      expect(handledAfterUnmount).toBe(false)
+      expect(onDelete).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('Exports', () => {
+    it('should export INSERT_CONTEXT_BLOCK_COMMAND', () => {
+      expect(INSERT_CONTEXT_BLOCK_COMMAND).toBeDefined()
+    })
+
+    it('should export DELETE_CONTEXT_BLOCK_COMMAND', () => {
+      expect(DELETE_CONTEXT_BLOCK_COMMAND).toBeDefined()
+    })
+
+    it('should export ContextBlock component', () => {
+      expect(ContextBlock).toBeDefined()
+    })
+  })
+
+  describe('Edge Cases', () => {
+    it('should handle undefined datasets prop', () => {
+      renderWithEditor(<ContextBlock datasets={undefined} />)
+
+      dispatchInsert()
+      expect(mockCreateContextBlockNode).toHaveBeenCalledWith([], expect.any(Function), undefined)
+    })
+
+    it('should handle empty datasets array', () => {
+      renderWithEditor(<ContextBlock datasets={[]} />)
+
+      dispatchInsert()
+      expect(mockCreateContextBlockNode).toHaveBeenCalledWith([], expect.any(Function), undefined)
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/context-block/node.spec.tsx

@@ -0,0 +1,244 @@
+import { $getRoot } from 'lexical'
+import { createTestEditor, withEditorUpdate } from '../__tests__/utils'
+import { $createContextBlockNode, $isContextBlockNode, ContextBlockNode } from './node'
+
+const mockDatasets = [
+  { id: '1', name: 'Dataset A', type: 'text' },
+  { id: '2', name: 'Dataset B', type: 'text' },
+]
+const mockOnAddContext = vi.fn()
+const createContextBlockTestEditor = () => createTestEditor([ContextBlockNode])
+describe('ContextBlockNode', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  describe('Static Methods', () => {
+    it('should return correct type', () => {
+      expect(ContextBlockNode.getType()).toBe('context-block')
+    })
+
+    it('should clone a node', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext, true)
+        $getRoot().append(node)
+        const cloned = ContextBlockNode.clone(node)
+        expect(cloned).toBeInstanceOf(ContextBlockNode)
+      })
+    })
+  })
+
+  describe('Constructor', () => {
+    it('should store datasets', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        $getRoot().append(node)
+        expect(node.getDatasets()).toEqual(mockDatasets)
+      })
+    })
+
+    it('should store onAddContext callback', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        $getRoot().append(node)
+        expect(node.getOnAddContext()).toBe(mockOnAddContext)
+      })
+    })
+
+    it('should store canNotAddContext', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext, true)
+        $getRoot().append(node)
+        expect(node.getCanNotAddContext()).toBe(true)
+      })
+    })
+
+    it('should default canNotAddContext to false', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        $getRoot().append(node)
+        expect(node.getCanNotAddContext()).toBe(false)
+      })
+    })
+  })
+
+  describe('isInline', () => {
+    it('should return true', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        expect(node.isInline()).toBe(true)
+      })
+    })
+  })
+
+  describe('createDOM', () => {
+    it('should create a div element', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        const dom = node.createDOM()
+        expect(dom.tagName).toBe('DIV')
+      })
+    })
+
+    it('should add correct CSS classes', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        const dom = node.createDOM()
+        expect(dom.classList.contains('inline-flex')).toBe(true)
+        expect(dom.classList.contains('items-center')).toBe(true)
+        expect(dom.classList.contains('align-middle')).toBe(true)
+      })
+    })
+  })
+
+  describe('updateDOM', () => {
+    it('should return false', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        expect(node.updateDOM()).toBe(false)
+      })
+    })
+  })
+
+  describe('decorate', () => {
+    it('should return a React element', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext, true)
+        $getRoot().append(node)
+        const result = node.decorate()
+        expect(result).toBeDefined()
+        expect(result.props).toEqual(
+          expect.objectContaining({
+            datasets: mockDatasets,
+            onAddContext: mockOnAddContext,
+            canNotAddContext: true,
+          }),
+        )
+      })
+    })
+
+    it('should pass nodeKey prop', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        $getRoot().append(node)
+        const result = node.decorate()
+        expect(result.props.nodeKey).toBe(node.getKey())
+      })
+    })
+  })
+
+  describe('getTextContent', () => {
+    it('should return the context placeholder', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        expect(node.getTextContent()).toBe('{{#context#}}')
+      })
+    })
+  })
+
+  describe('exportJSON', () => {
+    it('should export correct JSON structure', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext, true)
+        $getRoot().append(node)
+        const json = node.exportJSON()
+        expect(json.type).toBe('context-block')
+        expect(json.version).toBe(1)
+        expect(json.datasets).toEqual(mockDatasets)
+        expect(json.onAddContext).toBe(mockOnAddContext)
+        expect(json.canNotAddContext).toBe(true)
+      })
+    })
+  })
+
+  describe('importJSON', () => {
+    it('should create a node from serialized data', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const serialized = {
+          type: 'context-block' as const,
+          version: 1,
+          datasets: mockDatasets,
+          onAddContext: mockOnAddContext,
+          canNotAddContext: false,
+        }
+        const node = ContextBlockNode.importJSON(serialized)
+        $getRoot().append(node)
+        expect(node).toBeInstanceOf(ContextBlockNode)
+        expect(node.getDatasets()).toEqual(mockDatasets)
+        expect(node.getOnAddContext()).toBe(mockOnAddContext)
+        expect(node.getCanNotAddContext()).toBe(false)
+      })
+    })
+  })
+
+  describe('$createContextBlockNode', () => {
+    it('should create a ContextBlockNode instance', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        expect(node).toBeInstanceOf(ContextBlockNode)
+      })
+    })
+
+    it('should pass canNotAddContext when provided', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext, true)
+        $getRoot().append(node)
+        expect(node.getCanNotAddContext()).toBe(true)
+      })
+    })
+  })
+
+  describe('$isContextBlockNode', () => {
+    it('should return true for ContextBlockNode instances', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext)
+        expect($isContextBlockNode(node)).toBe(true)
+      })
+    })
+
+    it('should return false for null', () => {
+      expect($isContextBlockNode(null)).toBe(false)
+    })
+
+    it('should return false for undefined', () => {
+      expect($isContextBlockNode(undefined)).toBe(false)
+    })
+  })
+
+  describe('Edge Cases', () => {
+    it('should handle empty datasets', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode([], mockOnAddContext)
+        $getRoot().append(node)
+        expect(node.getDatasets()).toEqual([])
+      })
+    })
+
+    it('should handle canNotAddContext as false explicitly', () => {
+      const editor = createContextBlockTestEditor()
+      withEditorUpdate(editor, () => {
+        const node = $createContextBlockNode(mockDatasets, mockOnAddContext, false)
+        $getRoot().append(node)
+        expect(node.getCanNotAddContext()).toBe(false)
+      })
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/custom-text/node.spec.tsx

@@ -0,0 +1,141 @@
+import type { EditorConfig, LexicalEditor } from 'lexical'
+import { $createParagraphNode, $getRoot } from 'lexical'
+import { createTestEditor, withEditorUpdate } from '../__tests__/utils'
+import { $createCustomTextNode, CustomTextNode } from './node'
+
+const createCustomTextTestEditor = () => createTestEditor([CustomTextNode])
+
+describe('CustomTextNode', () => {
+  let editor: LexicalEditor
+
+  beforeEach(() => {
+    editor = createCustomTextTestEditor()
+  })
+
+  afterEach(() => {
+    editor.setRootElement(null)
+  })
+
+  describe('Static Methods', () => {
+    it('should return correct type', () => {
+      expect(CustomTextNode.getType()).toBe('custom-text')
+    })
+
+    it('should clone a node', () => {
+      withEditorUpdate(editor, () => {
+        const paragraph = $createParagraphNode()
+        $getRoot().append(paragraph)
+        const node = $createCustomTextNode('hello')
+        paragraph.append(node)
+        const cloned = CustomTextNode.clone(node)
+        expect(cloned).toBeInstanceOf(CustomTextNode)
+      })
+    })
+  })
+
+  describe('createDOM', () => {
+    it('should create a DOM element', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('test')
+        const config: EditorConfig = { namespace: 'test', theme: {} }
+        const dom = node.createDOM(config)
+        expect(dom).toBeDefined()
+      })
+    })
+  })
+
+  describe('exportJSON', () => {
+    it('should export correct JSON structure', () => {
+      withEditorUpdate(editor, () => {
+        const paragraph = $createParagraphNode()
+        $getRoot().append(paragraph)
+        const node = $createCustomTextNode('hello world')
+        paragraph.append(node)
+        const json = node.exportJSON()
+        expect(json.type).toBe('custom-text')
+        expect(json.version).toBe(1)
+        expect(json.text).toBe('hello world')
+        expect(json.format).toBeDefined()
+        expect(json.detail).toBeDefined()
+        expect(json.style).toBeDefined()
+      })
+    })
+  })
+
+  describe('importJSON', () => {
+    it('should create a text node from serialized data', () => {
+      withEditorUpdate(editor, () => {
+        const serialized = {
+          type: 'custom-text' as const,
+          version: 1,
+          text: 'imported text',
+          format: 0,
+          detail: 0,
+          mode: 'normal' as const,
+          style: '',
+        }
+        const node = CustomTextNode.importJSON(serialized)
+        expect(node).toBeDefined()
+        expect(node.getTextContent()).toBe('imported text')
+      })
+    })
+  })
+
+  describe('isSimpleText', () => {
+    it('should return true for custom-text type with mode 0', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('simple')
+        expect(node.isSimpleText()).toBe(true)
+      })
+    })
+  })
+
+  describe('getTextContent', () => {
+    it('should return the text content', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('my content')
+        expect(node.getTextContent()).toBe('my content')
+      })
+    })
+  })
+
+  describe('$createCustomTextNode', () => {
+    it('should create a CustomTextNode instance', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('test')
+        expect(node).toBeInstanceOf(CustomTextNode)
+      })
+    })
+
+    it('should set the text content', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('hello')
+        expect(node.getTextContent()).toBe('hello')
+      })
+    })
+  })
+
+  describe('Edge Cases', () => {
+    it('should handle empty string', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('')
+        expect(node.getTextContent()).toBe('')
+      })
+    })
+
+    it('should handle special characters', () => {
+      withEditorUpdate(editor, () => {
+        const node = $createCustomTextNode('{{#context#}}')
+        expect(node.getTextContent()).toBe('{{#context#}}')
+      })
+    })
+
+    it('should handle very long text', () => {
+      withEditorUpdate(editor, () => {
+        const longText = 'A'.repeat(10000)
+        const node = $createCustomTextNode(longText)
+        expect(node.getTextContent()).toBe(longText)
+      })
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/draggable-plugin/index.spec.tsx

@@ -0,0 +1,112 @@
+import { LexicalComposer } from '@lexical/react/LexicalComposer'
+import { ContentEditable } from '@lexical/react/LexicalContentEditable'
+import { LexicalErrorBoundary } from '@lexical/react/LexicalErrorBoundary'
+import { RichTextPlugin } from '@lexical/react/LexicalRichTextPlugin'
+import { fireEvent, render, screen, waitFor } from '@testing-library/react'
+import DraggableBlockPlugin from '.'
+
+const CONTENT_EDITABLE_TEST_ID = 'draggable-content-editable'
+let namespaceCounter = 0
+
+function renderWithEditor(anchorElem?: HTMLElement) {
+  render(
+    <LexicalComposer
+      initialConfig={{
+        namespace: `draggable-plugin-test-${namespaceCounter++}`,
+        onError: (error: Error) => { throw error },
+      }}
+    >
+      <RichTextPlugin
+        contentEditable={<ContentEditable data-testid={CONTENT_EDITABLE_TEST_ID} />}
+        placeholder={null}
+        ErrorBoundary={LexicalErrorBoundary}
+      />
+      <DraggableBlockPlugin anchorElem={anchorElem} />
+    </LexicalComposer>,
+  )
+
+  return screen.getByTestId(CONTENT_EDITABLE_TEST_ID)
+}
+
+function appendChildToRoot(rootElement: HTMLElement, className = '') {
+  const element = document.createElement('div')
+  element.className = className
+  rootElement.appendChild(element)
+  return element
+}
+
+describe('DraggableBlockPlugin', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  describe('Rendering', () => {
+    it('should use body as default anchor and render target line', () => {
+      renderWithEditor()
+
+      const targetLine = screen.getByTestId('draggable-target-line')
+      expect(targetLine).toBeInTheDocument()
+      expect(document.body.contains(targetLine)).toBe(true)
+      expect(screen.queryByTestId('draggable-menu')).not.toBeInTheDocument()
+    })
+
+    it('should render inside custom anchor element when provided', () => {
+      const customAnchor = document.createElement('div')
+      document.body.appendChild(customAnchor)
+
+      renderWithEditor(customAnchor)
+
+      const targetLine = screen.getByTestId('draggable-target-line')
+      expect(customAnchor.contains(targetLine)).toBe(true)
+
+      customAnchor.remove()
+    })
+  })
+
+  describe('Drag Support Detection', () => {
+    it('should render drag menu when mouse moves over a support-drag element', async () => {
+      const rootElement = renderWithEditor()
+      const supportDragTarget = appendChildToRoot(rootElement, 'support-drag')
+
+      expect(screen.queryByTestId('draggable-menu')).not.toBeInTheDocument()
+      fireEvent.mouseMove(supportDragTarget)
+
+      await waitFor(() => {
+        expect(screen.getByTestId('draggable-menu')).toBeInTheDocument()
+      })
+    })
+
+    it('should hide drag menu when support-drag target is removed and mouse moves again', async () => {
+      const rootElement = renderWithEditor()
+      const supportDragTarget = appendChildToRoot(rootElement, 'support-drag')
+
+      fireEvent.mouseMove(supportDragTarget)
+      await waitFor(() => {
+        expect(screen.getByTestId('draggable-menu')).toBeInTheDocument()
+      })
+
+      supportDragTarget.remove()
+      fireEvent.mouseMove(rootElement)
+      await waitFor(() => {
+        expect(screen.queryByTestId('draggable-menu')).not.toBeInTheDocument()
+      })
+    })
+  })
+
+  describe('Menu Detection Contract', () => {
+    it('should render menu with draggable-block-menu class and keep non-menu elements outside it', async () => {
+      const rootElement = renderWithEditor()
+      const supportDragTarget = appendChildToRoot(rootElement, 'support-drag')
+
+      fireEvent.mouseMove(supportDragTarget)
+
+      const menuIcon = await screen.findByTestId('draggable-menu-icon')
+      expect(menuIcon.closest('.draggable-block-menu')).not.toBeNull()
+
+      const normalElement = document.createElement('div')
+      document.body.appendChild(normalElement)
+      expect(normalElement.closest('.draggable-block-menu')).toBeNull()
+      normalElement.remove()
+    })
+  })
+})

web/app/components/base/prompt-editor/plugins/draggable-plugin/index.tsx

@@ -1,7 +1,6 @@
 import type { JSX } from 'react'
 import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
 import { DraggableBlockPlugin_EXPERIMENTAL } from '@lexical/react/LexicalDraggableBlockPlugin'
-import { RiDraggable } from '@remixicon/react'
 import { useEffect, useRef, useState } from 'react'
 import { cn } from '@/utils/classnames'
 
@@ -61,8 +60,8 @@ export default function DraggableBlockPlugin({
       menuComponent={
         isSupportDrag
           ? (
-              <div ref={menuRef} className={cn(DRAGGABLE_BLOCK_MENU_CLASSNAME, 'absolute right-2.5 top-4 cursor-grab opacity-0 will-change-transform active:cursor-move')}>
-                <RiDraggable className="size-3.5 text-text-tertiary" />
+              <div ref={menuRef} className={cn(DRAGGABLE_BLOCK_MENU_CLASSNAME, 'absolute right-2.5 top-4 cursor-grab opacity-0 will-change-transform active:cursor-move')} data-testid="draggable-menu">
+                <span className="i-ri-draggable size-3.5 text-text-tertiary" data-testid="draggable-menu-icon" />
               </div>
             )
           : null
@@ -71,6 +70,7 @@ export default function DraggableBlockPlugin({
         <div
           ref={targetLineRef}
           className="pointer-events-none absolute left-[-21px] top-0 opacity-0 will-change-transform"
+          data-testid="draggable-target-line"
           // style={{ width: 500 }} // width not worked here
         >
           <div

web/app/components/base/prompt-editor/utils.spec.ts

@@ -0,0 +1,267 @@
+import type {
+  Klass,
+  LexicalEditor,
+  LexicalNode,
+  RangeSelection,
+  TextNode,
+} from 'lexical'
+import type { CustomTextNode } from './plugins/custom-text/node'
+import type { MenuTextMatch } from './types'
+import {
+  $splitNodeContainingQuery,
+  decoratorTransform,
+  getSelectedNode,
+  registerLexicalTextEntity,
+  textToEditorState,
+} from './utils'
+
+const mockState = vi.hoisted(() => ({
+  isAtNodeEnd: false,
+  selection: null as unknown,
+  createTextNode: vi.fn(),
+}))
+
+vi.mock('@lexical/selection', () => ({
+  $isAtNodeEnd: () => mockState.isAtNodeEnd,
+}))
+
+vi.mock('lexical', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('lexical')>()
+  return {
+    ...actual,
+    $getSelection: () => mockState.selection,
+    $isRangeSelection: (selection: unknown) => !!(selection as { __isRangeSelection?: boolean } | null)?.__isRangeSelection,
+    $createTextNode: mockState.createTextNode,
+    $isTextNode: (node: unknown) => !!(node as { __isTextNode?: boolean } | null)?.__isTextNode,
+  }
+})
+
+vi.mock('./plugins/custom-text/node', () => ({
+  CustomTextNode: class MockCustomTextNode {},
+}))
+
+describe('prompt-editor/utils', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockState.isAtNodeEnd = false
+    mockState.selection = null
+  })
+
+  // Node selection utility for forward/backward lexical cursor behavior.
+  describe('getSelectedNode', () => {
+    it('should return anchor node when anchor and focus are the same node', () => {
+      const sharedNode = { id: 'same' }
+      const selection = {
+        anchor: { getNode: () => sharedNode },
+        focus: { getNode: () => sharedNode },
+        isBackward: () => false,
+      } as unknown as RangeSelection
+
+      expect(getSelectedNode(selection)).toBe(sharedNode)
+    })
+
+    it('should return anchor node for backward selection when focus is at node end', () => {
+      const anchorNode = { id: 'anchor' }
+      const focusNode = { id: 'focus' }
+      const selection = {
+        anchor: { getNode: () => anchorNode },
+        focus: { getNode: () => focusNode },
+        isBackward: () => true,
+      } as unknown as RangeSelection
+
+      mockState.isAtNodeEnd = true
+      expect(getSelectedNode(selection)).toBe(anchorNode)
+    })
+
+    it('should return focus node for forward selection when anchor is not at node end', () => {
+      const anchorNode = { id: 'anchor' }
+      const focusNode = { id: 'focus' }
+      const selection = {
+        anchor: { getNode: () => anchorNode },
+        focus: { getNode: () => focusNode },
+        isBackward: () => false,
+      } as unknown as RangeSelection
+
+      mockState.isAtNodeEnd = false
+      expect(getSelectedNode(selection)).toBe(focusNode)
+    })
+  })
+
+  // Entity registration should register transforms and convert invalid entity nodes.
+  describe('registerLexicalTextEntity', () => {
+    it('should register transforms and replace invalid target node with plain text', () => {
+      class TargetNode {
+        __isTextNode = true
+        getTextContent = vi.fn(() => 'invalid')
+        getFormat = vi.fn(() => 9)
+        replace = vi.fn()
+        splitText = vi.fn()
+        getPreviousSibling = vi.fn(() => null)
+        getNextSibling = vi.fn(() => null)
+        getLatest = vi.fn(() => ({ __mode: 0 }))
+      }
+
+      const removePlainTextTransform = vi.fn()
+      const removeReverseNodeTransform = vi.fn()
+      const registerNodeTransform = vi
+        .fn()
+        .mockReturnValueOnce(removePlainTextTransform)
+        .mockReturnValueOnce(removeReverseNodeTransform)
+      const editor = {
+        registerNodeTransform,
+      } as unknown as LexicalEditor
+      const createdTextNode = {
+        setFormat: vi.fn(),
+      }
+      mockState.createTextNode.mockReturnValue(createdTextNode)
+      const getMatch = vi.fn(() => null)
+      type TargetTextNode = InstanceType<typeof TargetNode> & TextNode
+      const targetNodeClass = TargetNode as unknown as Klass<TargetTextNode>
+      const createNode = vi.fn((textNode: TextNode) => textNode as TargetTextNode)
+
+      const cleanups = registerLexicalTextEntity(editor, getMatch, targetNodeClass, createNode)
+      expect(cleanups).toEqual([removePlainTextTransform, removeReverseNodeTransform])
+
+      const reverseNodeTransform = registerNodeTransform.mock.calls[1][1] as (node: TargetTextNode) => void
+      const targetNode = new TargetNode() as TargetTextNode
+      reverseNodeTransform(targetNode)
+
+      expect(mockState.createTextNode).toHaveBeenCalledWith('invalid')
+      expect(createdTextNode.setFormat).toHaveBeenCalledWith(9)
+      expect(targetNode.replace).toHaveBeenCalledWith(createdTextNode)
+    })
+  })
+
+  // Decorator transform behavior for converting matched text segments.
+  describe('decoratorTransform', () => {
+    it('should do nothing when node is not simple text', () => {
+      const node = {
+        isSimpleText: vi.fn(() => false),
+      } as unknown as CustomTextNode
+      const getMatch = vi.fn()
+      const createNode = vi.fn()
+
+      decoratorTransform(node, getMatch, createNode)
+
+      expect(getMatch).not.toHaveBeenCalled()
+      expect(createNode).not.toHaveBeenCalled()
+    })
+
+    it('should replace matched text node segment with created decorator node', () => {
+      const replacedNode = { replace: vi.fn() }
+      const node = {
+        __isTextNode: true,
+        isSimpleText: vi.fn(() => true),
+        getPreviousSibling: vi.fn(() => null),
+        getTextContent: vi.fn(() => 'abc'),
+        getNextSibling: vi.fn(() => null),
+        splitText: vi.fn(() => [replacedNode, null]),
+      } as unknown as CustomTextNode
+      const getMatch = vi
+        .fn()
+        .mockReturnValueOnce({ start: 0, end: 1 })
+        .mockReturnValueOnce(null)
+      const createdDecoratorNode = { id: 'decorator' }
+      const createNode = vi.fn(() => createdDecoratorNode as unknown as LexicalNode)
+
+      decoratorTransform(node, getMatch, createNode)
+
+      expect(node.splitText).toHaveBeenCalledWith(1)
+      expect(createNode).toHaveBeenCalledWith(replacedNode)
+      expect(replacedNode.replace).toHaveBeenCalledWith(createdDecoratorNode)
+    })
+  })
+
+  // Split helper for menu query replacement inside collapsed text selection.
+  describe('$splitNodeContainingQuery', () => {
+    const match: MenuTextMatch = {
+      leadOffset: 0,
+      matchingString: 'abc',
+      replaceableString: '@abc',
+    }
+
+    it('should return null when selection is not a collapsed range selection', () => {
+      mockState.selection = { __isRangeSelection: false }
+      expect($splitNodeContainingQuery(match)).toBeNull()
+    })
+
+    it('should return null when anchor is not text selection', () => {
+      mockState.selection = {
+        __isRangeSelection: true,
+        isCollapsed: () => true,
+        anchor: {
+          type: 'element',
+          offset: 1,
+          getNode: vi.fn(),
+        },
+      }
+
+      expect($splitNodeContainingQuery(match)).toBeNull()
+    })
+
+    it('should split using single offset when query starts at beginning of text', () => {
+      const newNode = { id: 'new-node' }
+      const anchorNode = {
+        isSimpleText: () => true,
+        getTextContent: () => '@abc',
+        splitText: vi.fn(() => [newNode]),
+      }
+      mockState.selection = {
+        __isRangeSelection: true,
+        isCollapsed: () => true,
+        anchor: {
+          type: 'text',
+          offset: 4,
+          getNode: () => anchorNode,
+        },
+      }
+
+      const result = $splitNodeContainingQuery(match)
+
+      expect(anchorNode.splitText).toHaveBeenCalledWith(4)
+      expect(result).toBe(newNode)
+    })
+
+    it('should split using range offsets when query is inside text', () => {
+      const newNode = { id: 'new-node' }
+      const anchorNode = {
+        isSimpleText: () => true,
+        getTextContent: () => 'hello @abc',
+        splitText: vi.fn(() => [null, newNode]),
+      }
+      mockState.selection = {
+        __isRangeSelection: true,
+        isCollapsed: () => true,
+        anchor: {
+          type: 'text',
+          offset: 10,
+          getNode: () => anchorNode,
+        },
+      }
+
+      const result = $splitNodeContainingQuery(match)
+
+      expect(anchorNode.splitText).toHaveBeenCalledWith(6, 10)
+      expect(result).toBe(newNode)
+    })
+  })
+
+  // Serialization utility for prompt text -> lexical editor state JSON.
+  describe('textToEditorState', () => {
+    it('should serialize multiline text into paragraph nodes', () => {
+      const state = JSON.parse(textToEditorState('line-1\nline-2'))
+
+      expect(state.root.children).toHaveLength(2)
+      expect(state.root.children[0].children[0].text).toBe('line-1')
+      expect(state.root.children[1].children[0].text).toBe('line-2')
+      expect(state.root.type).toBe('root')
+    })
+
+    it('should create one empty paragraph when text is empty', () => {
+      const state = JSON.parse(textToEditorState(''))
+
+      expect(state.root.children).toHaveLength(1)
+      expect(state.root.children[0].children[0].text).toBe('')
+    })
+  })
+})

web/eslint-suppressions.json

@@ -2325,9 +2325,6 @@
     }
   },
   "app/components/base/prompt-editor/plugins/context-block/component.tsx": {
-    "tailwindcss/no-duplicate-classes": {
-      "count": 1
-    },
     "ts/no-explicit-any": {
       "count": 1
     }