custom-edge.tsx 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import {
  2. memo,
  3. useCallback,
  4. useState,
  5. } from 'react'
  6. import { intersection } from 'lodash-es'
  7. import type { EdgeProps } from 'reactflow'
  8. import {
  9. BaseEdge,
  10. EdgeLabelRenderer,
  11. Position,
  12. getSimpleBezierPath,
  13. } from 'reactflow'
  14. import {
  15. useNodesExtraData,
  16. useNodesInteractions,
  17. } from './hooks'
  18. import BlockSelector from './block-selector'
  19. import type {
  20. Edge,
  21. OnSelectBlock,
  22. } from './types'
  23. const CustomEdge = ({
  24. id,
  25. data,
  26. source,
  27. sourceHandleId,
  28. target,
  29. targetHandleId,
  30. sourceX,
  31. sourceY,
  32. targetX,
  33. targetY,
  34. selected,
  35. }: EdgeProps) => {
  36. const [
  37. edgePath,
  38. labelX,
  39. labelY,
  40. ] = getSimpleBezierPath({
  41. sourceX: sourceX - 8,
  42. sourceY,
  43. sourcePosition: Position.Right,
  44. targetX: targetX + 8,
  45. targetY,
  46. targetPosition: Position.Left,
  47. })
  48. const [open, setOpen] = useState(false)
  49. const { handleNodeAdd } = useNodesInteractions()
  50. const nodesExtraData = useNodesExtraData()
  51. const availablePrevNodes = nodesExtraData[(data as Edge['data'])!.targetType]?.availablePrevNodes || []
  52. const availableNextNodes = nodesExtraData[(data as Edge['data'])!.sourceType]?.availableNextNodes || []
  53. const handleOpenChange = useCallback((v: boolean) => {
  54. setOpen(v)
  55. }, [])
  56. const handleInsert = useCallback<OnSelectBlock>((nodeType, toolDefaultValue) => {
  57. handleNodeAdd(
  58. {
  59. nodeType,
  60. toolDefaultValue,
  61. },
  62. {
  63. prevNodeId: source,
  64. prevNodeSourceHandle: sourceHandleId || 'source',
  65. nextNodeId: target,
  66. nextNodeTargetHandle: targetHandleId || 'target',
  67. },
  68. )
  69. }, [handleNodeAdd, source, sourceHandleId, target, targetHandleId])
  70. return (
  71. <>
  72. <BaseEdge
  73. id={id}
  74. path={edgePath}
  75. style={{
  76. stroke: (selected || data?._connectedNodeIsHovering || data?._runned) ? '#2970FF' : '#D0D5DD',
  77. strokeWidth: 2,
  78. }}
  79. />
  80. <EdgeLabelRenderer>
  81. <div
  82. className={`
  83. nopan nodrag hover:scale-125
  84. ${data?._hovering ? 'block' : 'hidden'}
  85. ${open && '!block'}
  86. `}
  87. style={{
  88. position: 'absolute',
  89. transform: `translate(-50%, -50%) translate(${labelX}px, ${labelY}px)`,
  90. pointerEvents: 'all',
  91. }}
  92. >
  93. <BlockSelector
  94. open={open}
  95. onOpenChange={handleOpenChange}
  96. asChild
  97. onSelect={handleInsert}
  98. availableBlocksTypes={intersection(availablePrevNodes, availableNextNodes)}
  99. triggerClassName={() => 'hover:scale-150 transition-all'}
  100. />
  101. </div>
  102. </EdgeLabelRenderer>
  103. </>
  104. )
  105. }
  106. export default memo(CustomEdge)