use-workflow-node-started.ts 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. import { useCallback } from 'react'
  2. import {
  3. useReactFlow,
  4. useStoreApi,
  5. } from 'reactflow'
  6. import produce from 'immer'
  7. import type { NodeStartedResponse } from '@/types/workflow'
  8. import { NodeRunningStatus } from '@/app/components/workflow/types'
  9. import { useWorkflowStore } from '@/app/components/workflow/store'
  10. export const useWorkflowNodeStarted = () => {
  11. const store = useStoreApi()
  12. const workflowStore = useWorkflowStore()
  13. const reactflow = useReactFlow()
  14. const handleWorkflowNodeStarted = useCallback((
  15. params: NodeStartedResponse,
  16. containerParams: {
  17. clientWidth: number,
  18. clientHeight: number,
  19. },
  20. ) => {
  21. const { data } = params
  22. const {
  23. workflowRunningData,
  24. setWorkflowRunningData,
  25. } = workflowStore.getState()
  26. const {
  27. getNodes,
  28. setNodes,
  29. edges,
  30. setEdges,
  31. transform,
  32. } = store.getState()
  33. const nodes = getNodes()
  34. setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
  35. draft.tracing!.push({
  36. ...data,
  37. status: NodeRunningStatus.Running,
  38. })
  39. }))
  40. const {
  41. setViewport,
  42. } = reactflow
  43. const currentNodeIndex = nodes.findIndex(node => node.id === data.node_id)
  44. const currentNode = nodes[currentNodeIndex]
  45. const position = currentNode.position
  46. const zoom = transform[2]
  47. if (!currentNode.parentId) {
  48. setViewport({
  49. x: (containerParams.clientWidth - 400 - currentNode.width! * zoom) / 2 - position.x * zoom,
  50. y: (containerParams.clientHeight - currentNode.height! * zoom) / 2 - position.y * zoom,
  51. zoom: transform[2],
  52. })
  53. }
  54. const newNodes = produce(nodes, (draft) => {
  55. draft[currentNodeIndex].data._runningStatus = NodeRunningStatus.Running
  56. draft[currentNodeIndex].data._waitingRun = false
  57. })
  58. setNodes(newNodes)
  59. const newEdges = produce(edges, (draft) => {
  60. const incomeEdges = draft.filter((edge) => {
  61. return edge.target === data.node_id
  62. })
  63. incomeEdges.forEach((edge) => {
  64. const incomeNode = nodes.find(node => node.id === edge.source)!
  65. if (!incomeNode || !('data' in incomeNode))
  66. return
  67. if (
  68. (!incomeNode.data._runningBranchId && edge.sourceHandle === 'source')
  69. || (incomeNode.data._runningBranchId && edge.sourceHandle === incomeNode.data._runningBranchId)
  70. ) {
  71. edge.data = {
  72. ...edge.data,
  73. _sourceRunningStatus: incomeNode.data._runningStatus,
  74. _targetRunningStatus: NodeRunningStatus.Running,
  75. _waitingRun: false,
  76. }
  77. }
  78. })
  79. })
  80. setEdges(newEdges)
  81. }, [workflowStore, store, reactflow])
  82. return {
  83. handleWorkflowNodeStarted,
  84. }
  85. }