node-handle.tsx 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186
  1. import type { MouseEvent } from 'react'
  2. import {
  3. memo,
  4. useCallback,
  5. useEffect,
  6. useState,
  7. } from 'react'
  8. import {
  9. Handle,
  10. Position,
  11. } from 'reactflow'
  12. import { BlockEnum } from '../../../types'
  13. import type { Node } from '../../../types'
  14. import BlockSelector from '../../../block-selector'
  15. import type { ToolDefaultValue } from '../../../block-selector/types'
  16. import {
  17. useNodesExtraData,
  18. useNodesInteractions,
  19. useNodesReadOnly,
  20. } from '../../../hooks'
  21. import { useStore } from '../../../store'
  22. type NodeHandleProps = {
  23. handleId: string
  24. handleClassName?: string
  25. nodeSelectorClassName?: string
  26. } & Pick<Node, 'id' | 'data'>
  27. export const NodeTargetHandle = memo(({
  28. id,
  29. data,
  30. handleId,
  31. handleClassName,
  32. nodeSelectorClassName,
  33. }: NodeHandleProps) => {
  34. const [open, setOpen] = useState(false)
  35. const { handleNodeAdd } = useNodesInteractions()
  36. const nodesExtraData = useNodesExtraData()
  37. const { getNodesReadOnly } = useNodesReadOnly()
  38. const connected = data._connectedTargetHandleIds?.includes(handleId)
  39. const availablePrevNodes = nodesExtraData[data.type].availablePrevNodes
  40. const isConnectable = !!availablePrevNodes.length
  41. const handleOpenChange = useCallback((v: boolean) => {
  42. setOpen(v)
  43. }, [])
  44. const handleHandleClick = useCallback((e: MouseEvent) => {
  45. e.stopPropagation()
  46. if (!connected)
  47. setOpen(v => !v)
  48. }, [connected])
  49. const handleSelect = useCallback((type: BlockEnum, toolDefaultValue?: ToolDefaultValue) => {
  50. handleNodeAdd(
  51. {
  52. nodeType: type,
  53. toolDefaultValue,
  54. },
  55. {
  56. nextNodeId: id,
  57. nextNodeTargetHandle: handleId,
  58. },
  59. )
  60. }, [handleNodeAdd, id, handleId])
  61. return (
  62. <>
  63. <Handle
  64. id={handleId}
  65. type='target'
  66. position={Position.Left}
  67. className={`
  68. !w-4 !h-4 !bg-transparent !rounded-none !outline-none !border-none z-[1]
  69. after:absolute after:w-0.5 after:h-2 after:left-1.5 after:top-1 after:bg-primary-500
  70. hover:scale-125 transition-all
  71. ${!connected && 'after:opacity-0'}
  72. ${data.type === BlockEnum.Start && 'opacity-0'}
  73. ${handleClassName}
  74. `}
  75. isConnectable={isConnectable}
  76. onClick={handleHandleClick}
  77. >
  78. {
  79. !connected && isConnectable && !data._isInvalidConnection && !getNodesReadOnly() && (
  80. <BlockSelector
  81. open={open}
  82. onOpenChange={handleOpenChange}
  83. onSelect={handleSelect}
  84. asChild
  85. placement='left'
  86. triggerClassName={open => `
  87. hidden absolute left-0 top-0 pointer-events-none
  88. ${nodeSelectorClassName}
  89. group-hover:!flex
  90. ${data.selected && '!flex'}
  91. ${open && '!flex'}
  92. `}
  93. availableBlocksTypes={availablePrevNodes}
  94. />
  95. )
  96. }
  97. </Handle>
  98. </>
  99. )
  100. })
  101. NodeTargetHandle.displayName = 'NodeTargetHandle'
  102. export const NodeSourceHandle = memo(({
  103. id,
  104. data,
  105. handleId,
  106. handleClassName,
  107. nodeSelectorClassName,
  108. }: NodeHandleProps) => {
  109. const notInitialWorkflow = useStore(s => s.notInitialWorkflow)
  110. const [open, setOpen] = useState(false)
  111. const { handleNodeAdd } = useNodesInteractions()
  112. const nodesExtraData = useNodesExtraData()
  113. const { getNodesReadOnly } = useNodesReadOnly()
  114. const availableNextNodes = nodesExtraData[data.type].availableNextNodes
  115. const isConnectable = !!availableNextNodes.length
  116. const connected = data._connectedSourceHandleIds?.includes(handleId)
  117. const handleOpenChange = useCallback((v: boolean) => {
  118. setOpen(v)
  119. }, [])
  120. const handleHandleClick = useCallback((e: MouseEvent) => {
  121. e.stopPropagation()
  122. if (!connected)
  123. setOpen(v => !v)
  124. }, [connected])
  125. const handleSelect = useCallback((type: BlockEnum, toolDefaultValue?: ToolDefaultValue) => {
  126. handleNodeAdd(
  127. {
  128. nodeType: type,
  129. toolDefaultValue,
  130. },
  131. {
  132. prevNodeId: id,
  133. prevNodeSourceHandle: handleId,
  134. },
  135. )
  136. }, [handleNodeAdd, id, handleId])
  137. useEffect(() => {
  138. if (notInitialWorkflow && data.type === BlockEnum.Start)
  139. setOpen(true)
  140. }, [notInitialWorkflow, data.type])
  141. return (
  142. <>
  143. <Handle
  144. id={handleId}
  145. type='source'
  146. position={Position.Right}
  147. className={`
  148. !w-4 !h-4 !bg-transparent !rounded-none !outline-none !border-none z-[1]
  149. after:absolute after:w-0.5 after:h-2 after:right-1.5 after:top-1 after:bg-primary-500
  150. hover:scale-125 transition-all
  151. ${!connected && 'after:opacity-0'}
  152. ${handleClassName}
  153. `}
  154. isConnectable={isConnectable}
  155. onClick={handleHandleClick}
  156. >
  157. {
  158. !connected && isConnectable && !data._isInvalidConnection && !getNodesReadOnly() && (
  159. <BlockSelector
  160. open={open}
  161. onOpenChange={handleOpenChange}
  162. onSelect={handleSelect}
  163. asChild
  164. triggerClassName={open => `
  165. hidden absolute top-0 left-0 pointer-events-none
  166. ${nodeSelectorClassName}
  167. group-hover:!flex
  168. ${data.selected && '!flex'}
  169. ${open && '!flex'}
  170. `}
  171. availableBlocksTypes={availableNextNodes}
  172. />
  173. )
  174. }
  175. </Handle>
  176. </>
  177. )
  178. })
  179. NodeSourceHandle.displayName = 'NodeSourceHandle'