fix: prompt-editor

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

@@ -19,6 +19,7 @@ import type {
 import { useBasicTypeaheadTriggerMatch } from '../../hooks'
 import { INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND } from '../workflow-variable-block'
 import { INSERT_VARIABLE_VALUE_BLOCK_COMMAND } from '../variable-block'
+import { $splitNodeContainingQuery } from '../../utils'
 import type { PromptOption } from './prompt-option'
 import PromptMenu from './prompt-menu'
 import VariableMenu from './variable-menu'
@@ -95,11 +96,17 @@ const ComponentPicker = ({
   )
 
   const handleSelectWorkflowVariable = useCallback((variables: string[]) => {
+    editor.update(() => {
+      const needRemove = $splitNodeContainingQuery(checkForTriggerMatch(triggerString, editor)!)
+      if (needRemove)
+        needRemove.remove()
+    })
+
     if (variables[1] === 'sys.query' || variables[1] === 'sys.files')
       editor.dispatchCommand(INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND, [variables[1]])
     else
       editor.dispatchCommand(INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND, variables)
-  }, [editor])
+  }, [editor, checkForTriggerMatch, triggerString])
 
   const renderMenu = useCallback<MenuRenderFn<PromptOption | VariableOption>>((
     anchorElementRef,

web/app/components/base/prompt-editor/plugins/workflow-variable-block/index.tsx

@@ -3,13 +3,12 @@ import {
   useEffect,
 } from 'react'
 import {
-  $getNodeByKey,
-  $getPreviousSelection,
   $insertNodes,
   COMMAND_PRIORITY_EDITOR,
   createCommand,
 } from 'lexical'
 import { mergeRegister } from '@lexical/utils'
+import { } from '@lexical/react/LexicalTypeaheadMenuPlugin'
 import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
 import type { WorkflowVariableBlockType } from '../../types'
 import {
@@ -42,13 +41,6 @@ const WorkflowVariableBlock = memo(({
         INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND,
         (variables: string[]) => {
           const workflowVariableBlockNode = $createWorkflowVariableBlockNode(variables, getWorkflowNode)
-          const prevNodeKey = ($getPreviousSelection() as any)?.anchor?.key
-
-          if (prevNodeKey) {
-            const prevNode = $getNodeByKey(prevNodeKey)
-
-            prevNode?.remove()
-          }
 
           $insertNodes([workflowVariableBlockNode])
           if (onInsert)

web/app/components/base/prompt-editor/types.ts

@@ -61,3 +61,9 @@ export type WorkflowVariableBlockType = {
   onInsert?: () => void
   onDelete?: () => void
 }
+
+export type MenuTextMatch = {
+  leadOffset: number
+  matchingString: string
+  replaceableString: string
+}

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

@@ -9,10 +9,13 @@ import type {
 } from 'lexical'
 import {
   $createTextNode,
+  $getSelection,
+  $isRangeSelection,
   $isTextNode,
 } from 'lexical'
 import type { EntityMatch } from '@lexical/text'
 import { CustomTextNode } from './plugins/custom-text/node'
+import type { MenuTextMatch } from './types'
 
 export function getSelectedNode(
   selection: RangeSelection,
@@ -249,6 +252,49 @@ export const decoratorTransform = (
   }
 }
 
+function getFullMatchOffset(
+  documentText: string,
+  entryText: string,
+  offset: number,
+): number {
+  let triggerOffset = offset
+  for (let i = triggerOffset; i <= entryText.length; i++) {
+    if (documentText.substr(-i) === entryText.substr(0, i))
+      triggerOffset = i
+  }
+  return triggerOffset
+}
+
+export function $splitNodeContainingQuery(match: MenuTextMatch): TextNode | null {
+  const selection = $getSelection()
+  if (!$isRangeSelection(selection) || !selection.isCollapsed())
+    return null
+  const anchor = selection.anchor
+  if (anchor.type !== 'text')
+    return null
+  const anchorNode = anchor.getNode()
+  if (!anchorNode.isSimpleText())
+    return null
+  const selectionOffset = anchor.offset
+  const textContent = anchorNode.getTextContent().slice(0, selectionOffset)
+  const characterOffset = match.replaceableString.length
+  const queryOffset = getFullMatchOffset(
+    textContent,
+    match.matchingString,
+    characterOffset,
+  )
+  const startOffset = selectionOffset - queryOffset
+  if (startOffset < 0)
+    return null
+  let newNode
+  if (startOffset === 0)
+    [newNode] = anchorNode.splitText(selectionOffset)
+  else
+    [, newNode] = anchorNode.splitText(startOffset, selectionOffset)
+
+  return newNode
+}
+
 export function textToEditorState(text: string) {
   const paragraph = text.split('\n')