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

import type { MenuRenderFn } from '@lexical/react/LexicalTypeaheadMenuPlugin'
import type { LexicalEditor, TextNode } from 'lexical'
import type {
  AgentOutputBlockType,
  ContextBlockType,
  CurrentBlockType,
  ErrorMessageBlockType,
  ExternalToolBlockType,
  HistoryBlockType,
  LastRunBlockType,
  MenuTextMatch,
  QueryBlockType,
  RequestURLBlockType,
  VariableBlockType,
  WorkflowVariableBlockType,
} from '../../types'
import type { PickerBlockMenuOption } from './menu'
import type { EventEmitterValue } from '@/context/event-emitter'
import { flip, offset, shift, useFloating } from '@floating-ui/react'
import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
import { LexicalTypeaheadMenuPlugin } from '@lexical/react/LexicalTypeaheadMenuPlugin'
import { mergeRegister } from '@lexical/utils'
import { BLUR_COMMAND, COMMAND_PRIORITY_EDITOR, FOCUS_COMMAND, KEY_ESCAPE_COMMAND } from 'lexical'
import { Fragment, memo, useCallback, useEffect, useRef, useState } from 'react'
import ReactDOM from 'react-dom'
import { useTranslation } from 'react-i18next'
import { GeneratorType } from '@/app/components/app/configuration/config/automatic/types'
import VarReferenceVars, {
  VAR_REFERENCE_CHILD_POPUP_CLASS_NAME,
} from '@/app/components/workflow/nodes/_base/components/variable/var-reference-vars'
import { useEventEmitterContextContext } from '@/context/event-emitter'
import { useBasicTypeaheadTriggerMatch } from '../../hooks'
import { $splitNodeContainingQuery } from '../../utils'
import { INSERT_AGENT_OUTPUT_BLOCK_COMMAND } from '../agent-output-block/commands'
import { INSERT_CURRENT_BLOCK_COMMAND } from '../current-block'
import { INSERT_ERROR_MESSAGE_BLOCK_COMMAND } from '../error-message-block'
import { INSERT_LAST_RUN_BLOCK_COMMAND } from '../last-run-block'
import { INSERT_VARIABLE_VALUE_BLOCK_COMMAND } from '../variable-block'
import { INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND } from '../workflow-variable-block'
import { useOptions } from './hooks'

type ComponentPickerProps = {
  triggerString: string
  contextBlock?: ContextBlockType
  queryBlock?: QueryBlockType
  requestURLBlock?: RequestURLBlockType
  historyBlock?: HistoryBlockType
  variableBlock?: VariableBlockType
  externalToolBlock?: ExternalToolBlockType
  workflowVariableBlock?: WorkflowVariableBlockType
  agentOutputBlock?: AgentOutputBlockType
  currentBlock?: CurrentBlockType
  errorMessageBlock?: ErrorMessageBlockType
  lastRunBlock?: LastRunBlockType
  isSupportFileVar?: boolean
}
const ComponentPicker = ({
  triggerString,
  contextBlock,
  queryBlock,
  requestURLBlock,
  historyBlock,
  variableBlock,
  externalToolBlock,
  workflowVariableBlock,
  agentOutputBlock,
  currentBlock,
  errorMessageBlock,
  lastRunBlock,
  isSupportFileVar,
}: ComponentPickerProps) => {
  const { t } = useTranslation()
  const { eventEmitter } = useEventEmitterContextContext()
  const { refs, floatingStyles, isPositioned } = useFloating({
    placement: 'bottom-start',
    middleware: [
      offset(0), // fix hide cursor
      shift({
        padding: 8,
      }),
      flip(),
    ],
  })
  const [editor] = useLexicalComposerContext()
  const triggerMatchRef = useRef<MenuTextMatch | null>(null)
  const baseCheckForTriggerMatch = useBasicTypeaheadTriggerMatch(triggerString, {
    minLength: 0,
    maxLength: 75,
  })
  const checkForTriggerMatch = useCallback(
    (text: string, editor: LexicalEditor) => {
      const match = baseCheckForTriggerMatch(text, editor)
      triggerMatchRef.current = match
      return match
    },
    [baseCheckForTriggerMatch],
  )

  const [queryString, setQueryString] = useState<string | null>(null)
  const [blurHidden, setBlurHidden] = useState(false)
  const blurTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const showAgentOutputAction = triggerString === '/' && agentOutputBlock?.show

  const clearBlurTimer = useCallback(() => {
    if (blurTimerRef.current) {
      clearTimeout(blurTimerRef.current)
      blurTimerRef.current = null
    }
  }, [])

  useEffect(() => {
    const unregister = mergeRegister(
      editor.registerCommand(
        BLUR_COMMAND,
        (event) => {
          clearBlurTimer()
          const target = event?.relatedTarget as HTMLElement
          const isVariableMenuTarget =
            target?.classList?.contains('var-search-input') ||
            target?.closest?.(`.${VAR_REFERENCE_CHILD_POPUP_CLASS_NAME}`)
          if (!isVariableMenuTarget)
            blurTimerRef.current = setTimeout(() => setBlurHidden(true), 200)
          return false
        },
        COMMAND_PRIORITY_EDITOR,
      ),
      editor.registerCommand(
        FOCUS_COMMAND,
        () => {
          clearBlurTimer()
          setBlurHidden(false)
          return false
        },
        COMMAND_PRIORITY_EDITOR,
      ),
    )

    return () => {
      if (blurTimerRef.current) clearTimeout(blurTimerRef.current)
      unregister()
    }
  }, [editor, clearBlurTimer])

  eventEmitter?.useSubscription((v: EventEmitterValue) => {
    if (
      typeof v !== 'string' &&
      v.type === INSERT_VARIABLE_VALUE_BLOCK_COMMAND &&
      typeof v.payload === 'string'
    )
      editor.dispatchCommand(INSERT_VARIABLE_VALUE_BLOCK_COMMAND, `{{${v.payload}}}`)
  })

  const { allFlattenOptions, workflowVariableOptions } = useOptions(
    contextBlock,
    queryBlock,
    historyBlock,
    variableBlock,
    externalToolBlock,
    workflowVariableBlock,
    requestURLBlock,
    currentBlock,
    errorMessageBlock,
    lastRunBlock,
    queryString || undefined,
  )

  const onSelectOption = useCallback(
    (
      selectedOption: PickerBlockMenuOption,
      nodeToRemove: TextNode | null,
      closeMenu: () => void,
    ) => {
      editor.update(() => {
        if (nodeToRemove && selectedOption?.key) nodeToRemove.remove()

        selectedOption.onSelectMenuOption()
        closeMenu()
      })
    },
    [editor],
  )

  const handleSelectWorkflowVariable = useCallback(
    (variables: string[]) => {
      editor.update(() => {
        const currentTriggerMatch =
          triggerMatchRef.current ?? checkForTriggerMatch(triggerString, editor)
        const needRemove = currentTriggerMatch
          ? $splitNodeContainingQuery(currentTriggerMatch)
          : null
        if (needRemove) needRemove.remove()
      })
      const isFlat = variables.length === 1
      if (isFlat) {
        const varName = variables[0]
        if (varName === 'current')
          editor.dispatchCommand(INSERT_CURRENT_BLOCK_COMMAND, currentBlock?.generatorType)
        else if (varName === 'error_message')
          editor.dispatchCommand(INSERT_ERROR_MESSAGE_BLOCK_COMMAND, null)
        else if (varName === 'last_run') editor.dispatchCommand(INSERT_LAST_RUN_BLOCK_COMMAND, null)
      } else 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, currentBlock?.generatorType, checkForTriggerMatch, triggerString],
  )

  const resetTypeaheadState = useCallback(() => {
    triggerMatchRef.current = null
    setQueryString(null)
    setBlurHidden(true)
  }, [])

  const handleClose = useCallback(() => {
    const escapeEvent = new KeyboardEvent('keydown', { key: 'Escape' })
    editor.dispatchCommand(KEY_ESCAPE_COMMAND, escapeEvent)
  }, [editor])

  const renderMenu = useCallback<MenuRenderFn<PickerBlockMenuOption>>(
    (anchorElementRef, { options, selectedIndex, selectOptionAndCleanUp, setHighlightedIndex }) => {
      const effectiveQueryString = triggerMatchRef.current?.matchingString ?? queryString

      if (blurHidden) return null
      if (
        !(
          anchorElementRef.current &&
          (allFlattenOptions.length || workflowVariableBlock?.show || showAgentOutputAction)
        )
      )
        return null

      setTimeout(() => {
        if (anchorElementRef.current) refs.setReference(anchorElementRef.current)
      }, 0)

      return (
        <>
          {ReactDOM.createPortal(
            // The `LexicalMenu` will try to calculate the position of the floating menu based on the first child.
            // Since we use floating ui, we need to wrap it with a div to prevent the position calculation being affected.
            // See https://github.com/facebook/lexical/blob/ac97dfa9e14a73ea2d6934ff566282d7f758e8bb/packages/lexical-react/src/shared/LexicalMenu.ts#L493
            <div className="size-0">
              <div
                className="w-[260px] rounded-lg border-[0.5px] border-components-panel-border bg-components-panel-bg-blur p-1 shadow-lg"
                style={{
                  ...floatingStyles,
                  visibility: isPositioned ? 'visible' : 'hidden',
                }}
                ref={refs.setFloating}
              >
                {workflowVariableBlock?.show && (
                  <div className="p-1">
                    <VarReferenceVars
                      hideSearch={triggerString === '/'}
                      searchText={triggerString === '/' ? effectiveQueryString || '' : undefined}
                      searchBoxClassName="mt-1"
                      vars={workflowVariableOptions}
                      onChange={(variables: string[]) => {
                        handleSelectWorkflowVariable(variables)
                      }}
                      maxHeightClass="max-h-[34vh]"
                      isSupportFileVar={isSupportFileVar}
                      onClose={handleClose}
                      onBlur={handleClose}
                      showManageInputField={workflowVariableBlock.showManageInputField}
                      onManageInputField={workflowVariableBlock.onManageInputField}
                      autoFocus={false}
                      isInCodeGeneratorInstructionEditor={
                        currentBlock?.generatorType === GeneratorType.code
                      }
                    />
                  </div>
                )}
                {workflowVariableBlock?.show && !!options.length && (
                  <div className="my-1 h-px w-full -translate-x-1 bg-divider-subtle"></div>
                )}
                <div>
                  {options.map((option, index) => (
                    <Fragment key={option.key}>
                      {
                        // Divider
                        index !== 0 && options.at(index - 1)?.group !== option.group && (
                          <div className="my-1 h-px w-full -translate-x-1 bg-divider-subtle"></div>
                        )
                      }
                      {option.renderMenuOption({
                        queryString: effectiveQueryString,
                        isSelected: workflowVariableBlock?.show ? false : selectedIndex === index,
                        onSelect: () => {
                          selectOptionAndCleanUp(option)
                        },
                        onSetHighlight: () => {
                          setHighlightedIndex(index)
                        },
                      })}
                    </Fragment>
                  ))}
                </div>
                {showAgentOutputAction && (
                  <div className="mt-1 border-t border-divider-subtle p-1">
                    <button
                      type="button"
                      className="flex h-6 w-full items-center gap-1 rounded-md py-1 pr-1 pl-3 text-left system-sm-regular text-text-secondary hover:bg-state-base-hover focus-visible:bg-state-base-hover focus-visible:ring-2 focus-visible:ring-state-accent-solid focus-visible:outline-hidden"
                      onMouseDown={(event) => event.preventDefault()}
                      onClick={() => {
                        editor.update(() => {
                          const currentTriggerMatch =
                            triggerMatchRef.current ?? checkForTriggerMatch(triggerString, editor)
                          const needRemove = currentTriggerMatch
                            ? $splitNodeContainingQuery(currentTriggerMatch)
                            : null
                          if (needRemove) needRemove.remove()
                        })
                        editor.dispatchCommand(INSERT_AGENT_OUTPUT_BLOCK_COMMAND, undefined)
                        resetTypeaheadState()
                      }}
                    >
                      <span aria-hidden="true" className="i-ri-add-line size-4 shrink-0" />
                      <span className="min-w-0 flex-1 truncate">
                        {t(($) => $['nodes.agent.outputVars.newOutput'], { ns: 'workflow' })}
                      </span>
                      <span
                        aria-hidden="true"
                        className="i-ri-question-line size-3.5 shrink-0 text-text-quaternary"
                      />
                    </button>
                  </div>
                )}
              </div>
            </div>,
            anchorElementRef.current,
          )}
        </>
      )
    },
    [
      blurHidden,
      allFlattenOptions.length,
      workflowVariableBlock?.show,
      showAgentOutputAction,
      floatingStyles,
      isPositioned,
      refs,
      workflowVariableOptions,
      isSupportFileVar,
      handleClose,
      currentBlock?.generatorType,
      handleSelectWorkflowVariable,
      queryString,
      triggerString,
      workflowVariableBlock?.showManageInputField,
      workflowVariableBlock?.onManageInputField,
      editor,
      checkForTriggerMatch,
      t,
      resetTypeaheadState,
    ],
  )

  return (
    <LexicalTypeaheadMenuPlugin
      options={allFlattenOptions}
      onQueryChange={setQueryString}
      onSelectOption={onSelectOption}
      // The `translate` class is used to workaround the issue that the `typeahead-menu` menu is not positioned as expected.
      // See also https://github.com/facebook/lexical/blob/772520509308e8ba7e4a82b6cd1996a78b3298d0/packages/lexical-react/src/shared/LexicalMenu.ts#L498
      //
      // We no need the position function of the `LexicalTypeaheadMenuPlugin`,
      // so the reference anchor should be positioned based on the range of the trigger string, and the menu will be positioned by the floating ui.
      anchorClassName="z-50 translate-y-[calc(-100%-3px)]"
      menuRenderFn={renderMenu}
      triggerFn={checkForTriggerMatch}
    />
  )
}

export default memo(ComponentPicker)