472 lines
16 KiB
TypeScript
472 lines
16 KiB
TypeScript
import { useCallback, useEffect, useMemo, useRef, useState, type MouseEvent as ReactMouseEvent } from 'react'
|
|
import { addEdge, useEdgesState, useNodesState, type Connection, type Edge, type Node, type ReactFlowInstance } from '@xyflow/react'
|
|
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
|
|
import { toast } from 'sonner'
|
|
|
|
import {
|
|
dispatchWorkflowDraft,
|
|
getNodeDefinitions,
|
|
getWorkflowRunComparison,
|
|
getWorkflowRuns,
|
|
preflightWorkflowDraft,
|
|
type WorkflowConfig,
|
|
type WorkflowDefinition,
|
|
type WorkflowExecutionMode,
|
|
type WorkflowNodeDefinition,
|
|
type WorkflowParams,
|
|
type WorkflowPreflightResponse,
|
|
} from '../../api/workflows'
|
|
import {
|
|
applyAutoLayout,
|
|
buildCurrentWorkflowConfig,
|
|
inferNodeLabel,
|
|
inferNodeType,
|
|
inferStepFromNodeType,
|
|
normalizeWorkflowParams,
|
|
resolveParamsForStepChange,
|
|
type WorkflowCanvasNodeData,
|
|
validateWorkflowDraft,
|
|
workflowToGraph,
|
|
} from './workflowGraphDraft'
|
|
import {
|
|
inferWorkflowFamily,
|
|
} from './workflowBlueprints'
|
|
import {
|
|
GRAPH_FAMILY_LABELS,
|
|
isDefinitionAllowedForGraphFamily,
|
|
} from './workflowNodeLibrary'
|
|
import type { WorkflowUtilityTab } from './WorkflowCanvasUtilitySidebar'
|
|
|
|
export type NodeMenuAnchor = {
|
|
clientX: number
|
|
clientY: number
|
|
flowPosition: { x: number; y: number }
|
|
}
|
|
|
|
function buildNodeData(
|
|
step: string,
|
|
params: WorkflowParams = {},
|
|
definition?: WorkflowNodeDefinition,
|
|
overrides?: Partial<WorkflowCanvasNodeData>,
|
|
): WorkflowCanvasNodeData {
|
|
return {
|
|
label: overrides?.label ?? definition?.label ?? inferNodeLabel(step),
|
|
params: normalizeWorkflowParams(params),
|
|
step,
|
|
description: overrides?.description ?? definition?.description,
|
|
icon: overrides?.icon ?? definition?.icon,
|
|
category: overrides?.category ?? definition?.category,
|
|
}
|
|
}
|
|
|
|
type UseWorkflowCanvasControllerArgs = {
|
|
workflow: WorkflowDefinition
|
|
onSave: (config: WorkflowConfig) => void
|
|
}
|
|
|
|
export function useWorkflowCanvasController({ workflow, onSave }: UseWorkflowCanvasControllerArgs) {
|
|
const queryClient = useQueryClient()
|
|
const { data: nodeDefinitionsData } = useQuery({
|
|
queryKey: ['workflow-node-definitions'],
|
|
queryFn: getNodeDefinitions,
|
|
staleTime: 5 * 60 * 1000,
|
|
})
|
|
const nodeDefinitions = nodeDefinitionsData?.definitions ?? []
|
|
const nodeDefinitionsByStep = Object.fromEntries(nodeDefinitions.map(definition => [definition.step, definition]))
|
|
|
|
const { nodes: initNodes, edges: initEdges } = workflowToGraph(workflow.config, nodeDefinitionsByStep)
|
|
const [nodes, setNodes, onNodesChange] = useNodesState(initNodes)
|
|
const [edges, setEdges, onEdgesChange] = useEdgesState(initEdges)
|
|
const [selectedNodeId, setSelectedNodeId] = useState<string | null>(null)
|
|
const [selectedRunId, setSelectedRunId] = useState<string | null>(null)
|
|
const [dispatchContextId, setDispatchContextId] = useState('')
|
|
const [preflightResult, setPreflightResult] = useState<WorkflowPreflightResponse | null>(null)
|
|
const [executionMode, setExecutionMode] = useState<WorkflowExecutionMode>(workflow.config.ui?.execution_mode ?? 'legacy')
|
|
const [nodeMenuAnchor, setNodeMenuAnchor] = useState<NodeMenuAnchor | null>(null)
|
|
const [activeUtilityTab, setActiveUtilityTab] = useState<WorkflowUtilityTab>('library')
|
|
const reactFlowWrapper = useRef<HTMLDivElement>(null)
|
|
const [reactFlowInstance, setReactFlowInstance] = useState<ReactFlowInstance<Node, Edge> | null>(null)
|
|
|
|
const validation = validateWorkflowDraft(nodes, edges, nodeDefinitionsByStep, nodeDefinitions.length > 0)
|
|
const selectedEdgeIds = useMemo(
|
|
() => edges.filter(edge => Boolean((edge as Edge & { selected?: boolean }).selected)).map(edge => edge.id),
|
|
[edges],
|
|
)
|
|
const graphFamily = useMemo(
|
|
() =>
|
|
inferWorkflowFamily(
|
|
buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode),
|
|
nodeDefinitionsByStep,
|
|
),
|
|
[edges, executionMode, nodeDefinitionsByStep, nodes, workflow],
|
|
)
|
|
|
|
const { data: workflowRuns = [] } = useQuery({
|
|
queryKey: ['workflow-runs', workflow.id],
|
|
queryFn: () => getWorkflowRuns(workflow.id),
|
|
refetchInterval: 5000,
|
|
})
|
|
|
|
const selectedRun = workflowRuns.find(run => run.id === selectedRunId) ?? workflowRuns[0] ?? null
|
|
|
|
const { data: selectedRunComparison, isFetching: isComparisonLoading } = useQuery({
|
|
queryKey: ['workflow-run-comparison', selectedRun?.id],
|
|
queryFn: () => getWorkflowRunComparison(selectedRun!.id),
|
|
enabled: Boolean(selectedRun?.id && selectedRun.execution_mode === 'shadow'),
|
|
refetchInterval: selectedRun?.status === 'pending' || selectedRun?.status === 'running' ? 5000 : false,
|
|
})
|
|
|
|
const dispatchMutation = useMutation({
|
|
mutationFn: ({ contextId, config }: { contextId: string; config: WorkflowConfig }) =>
|
|
dispatchWorkflowDraft({
|
|
workflow_id: workflow.id,
|
|
context_id: contextId,
|
|
config,
|
|
}),
|
|
onSuccess: result => {
|
|
queryClient.invalidateQueries({ queryKey: ['workflow-runs', workflow.id] })
|
|
setSelectedRunId(result.workflow_run.id)
|
|
toast.success(`Graph run dispatched: ${result.dispatched} task${result.dispatched === 1 ? '' : 's'}`)
|
|
},
|
|
onError: (error: any) => {
|
|
toast.error(error?.response?.data?.detail || 'Failed to dispatch workflow')
|
|
},
|
|
})
|
|
|
|
const preflightMutation = useMutation({
|
|
mutationFn: ({ contextId, config }: { contextId: string; config: WorkflowConfig }) =>
|
|
preflightWorkflowDraft({
|
|
workflow_id: workflow.id,
|
|
context_id: contextId,
|
|
config,
|
|
}),
|
|
onSuccess: result => {
|
|
setPreflightResult(result)
|
|
if (result.graph_dispatch_allowed) {
|
|
toast.success(result.summary)
|
|
} else {
|
|
toast.error(result.summary)
|
|
}
|
|
},
|
|
onError: (error: any) => {
|
|
setPreflightResult(null)
|
|
toast.error(error?.response?.data?.detail || 'Failed to preflight workflow')
|
|
},
|
|
})
|
|
|
|
useEffect(() => {
|
|
const graph = workflowToGraph(workflow.config, nodeDefinitionsByStep)
|
|
setNodes(graph.nodes)
|
|
setEdges(graph.edges)
|
|
setSelectedNodeId(null)
|
|
setSelectedRunId(null)
|
|
setNodeMenuAnchor(null)
|
|
setPreflightResult(null)
|
|
setExecutionMode(workflow.config.ui?.execution_mode ?? 'legacy')
|
|
setActiveUtilityTab('library')
|
|
}, [nodeDefinitionsData, setEdges, setNodes, workflow.config])
|
|
|
|
useEffect(() => {
|
|
if (!selectedRunId && workflowRuns.length > 0) {
|
|
setSelectedRunId(workflowRuns[0].id)
|
|
return
|
|
}
|
|
if (selectedRunId && !workflowRuns.some(run => run.id === selectedRunId)) {
|
|
setSelectedRunId(workflowRuns[0]?.id ?? null)
|
|
}
|
|
}, [selectedRunId, workflowRuns])
|
|
|
|
const onConnect = useCallback(
|
|
(connection: Connection) => setEdges(currentEdges => addEdge(connection, currentEdges)),
|
|
[setEdges],
|
|
)
|
|
|
|
const onNodeClick = useCallback((_: ReactMouseEvent, node: Node) => {
|
|
setNodeMenuAnchor(null)
|
|
setSelectedNodeId(node.id)
|
|
setActiveUtilityTab('inspector')
|
|
}, [])
|
|
|
|
const onEdgeClick = useCallback((_: ReactMouseEvent, edge: Edge) => {
|
|
setNodeMenuAnchor(null)
|
|
setSelectedNodeId(null)
|
|
setEdges(currentEdges =>
|
|
currentEdges.map(currentEdge => ({
|
|
...currentEdge,
|
|
selected: currentEdge.id === edge.id,
|
|
})),
|
|
)
|
|
}, [setEdges])
|
|
|
|
const onPaneClick = useCallback(() => {
|
|
setNodeMenuAnchor(null)
|
|
setSelectedNodeId(null)
|
|
setEdges(currentEdges =>
|
|
currentEdges.map(edge => ({
|
|
...edge,
|
|
selected: false,
|
|
})),
|
|
)
|
|
}, [setEdges])
|
|
|
|
const handleParamsChange = useCallback(
|
|
(newParams: WorkflowParams) => {
|
|
setNodes(currentNodes =>
|
|
currentNodes.map(node => {
|
|
if (node.id === selectedNodeId) {
|
|
return { ...node, data: { ...node.data, params: normalizeWorkflowParams(newParams) } }
|
|
}
|
|
return node
|
|
}),
|
|
)
|
|
},
|
|
[selectedNodeId, setNodes],
|
|
)
|
|
|
|
const handlePipelineStepChange = useCallback(
|
|
(stepName: string) => {
|
|
const definition = nodeDefinitionsByStep[stepName]
|
|
if (definition && !isDefinitionAllowedForGraphFamily(definition, graphFamily)) {
|
|
toast.error(`${definition.label} does not belong to the ${GRAPH_FAMILY_LABELS[graphFamily]} family.`)
|
|
return
|
|
}
|
|
|
|
setNodes(currentNodes =>
|
|
currentNodes.map(node => {
|
|
if (node.id !== selectedNodeId) return node
|
|
|
|
const currentData = (node.data as WorkflowCanvasNodeData | undefined) ?? buildNodeData(stepName)
|
|
return {
|
|
...node,
|
|
type: definition?.node_type ?? inferNodeType(stepName),
|
|
data: {
|
|
...buildNodeData(
|
|
stepName || inferStepFromNodeType(node.type),
|
|
resolveParamsForStepChange(definition, currentData.params),
|
|
definition,
|
|
),
|
|
step: stepName || inferStepFromNodeType(node.type),
|
|
},
|
|
}
|
|
}),
|
|
)
|
|
},
|
|
[graphFamily, nodeDefinitionsByStep, selectedNodeId, setNodes],
|
|
)
|
|
|
|
const openNodeMenu = useCallback(
|
|
(clientX: number, clientY: number) => {
|
|
if (!reactFlowInstance) return
|
|
setNodeMenuAnchor({
|
|
clientX,
|
|
clientY,
|
|
flowPosition: reactFlowInstance.screenToFlowPosition({ x: clientX, y: clientY }),
|
|
})
|
|
},
|
|
[reactFlowInstance],
|
|
)
|
|
|
|
const handlePaneContextMenu = useCallback(
|
|
(event: MouseEvent | ReactMouseEvent) => {
|
|
event.preventDefault()
|
|
setSelectedNodeId(null)
|
|
openNodeMenu(event.clientX, event.clientY)
|
|
},
|
|
[openNodeMenu],
|
|
)
|
|
|
|
const handleNodeContextMenu = useCallback(
|
|
(event: ReactMouseEvent, node: Node) => {
|
|
event.preventDefault()
|
|
setSelectedNodeId(node.id)
|
|
openNodeMenu(event.clientX, event.clientY)
|
|
},
|
|
[openNodeMenu],
|
|
)
|
|
|
|
const insertNode = useCallback(
|
|
(step: string, preferredPosition?: { x: number; y: number }) => {
|
|
const definition = nodeDefinitionsByStep[step]
|
|
if (definition && !isDefinitionAllowedForGraphFamily(definition, graphFamily)) {
|
|
toast.error(`${definition.label} cannot be added to a ${GRAPH_FAMILY_LABELS[graphFamily]} workflow.`)
|
|
return
|
|
}
|
|
|
|
const type = definition?.node_type ?? inferNodeType(step)
|
|
const fallbackX = nodes.length > 0 ? Math.max(...nodes.map(node => node.position.x)) + 220 : 120
|
|
const fallbackY = nodes.length > 0 ? Math.max(...nodes.map(node => node.position.y)) + 40 : 120
|
|
const newNode: Node = {
|
|
id: `${step}_${Date.now()}`,
|
|
type,
|
|
position: preferredPosition ?? { x: fallbackX, y: fallbackY },
|
|
data: buildNodeData(step, definition?.defaults ?? {}, definition),
|
|
}
|
|
|
|
setNodes(currentNodes => [...currentNodes, newNode])
|
|
setSelectedNodeId(newNode.id)
|
|
setNodeMenuAnchor(null)
|
|
setActiveUtilityTab('inspector')
|
|
},
|
|
[graphFamily, nodeDefinitionsByStep, nodes, setNodes],
|
|
)
|
|
|
|
const handleOpenToolbarNodeMenu = useCallback(() => {
|
|
if (!reactFlowWrapper.current || !reactFlowInstance) return
|
|
const bounds = reactFlowWrapper.current.getBoundingClientRect()
|
|
openNodeMenu(bounds.left + 36, bounds.top + 36)
|
|
}, [openNodeMenu, reactFlowInstance])
|
|
|
|
const handleAutoLayout = useCallback(() => {
|
|
setNodes(currentNodes => applyAutoLayout(currentNodes, edges))
|
|
setNodeMenuAnchor(null)
|
|
window.requestAnimationFrame(() => {
|
|
reactFlowInstance?.fitView({ padding: 0.2, duration: 250 })
|
|
})
|
|
}, [edges, reactFlowInstance, setNodes])
|
|
|
|
const deleteEdgesById = useCallback((edgeIds: string[]) => {
|
|
if (edgeIds.length === 0) return
|
|
setEdges(currentEdges => currentEdges.filter(edge => !edgeIds.includes(edge.id)))
|
|
setSelectedNodeId(null)
|
|
setNodeMenuAnchor(null)
|
|
toast.success(edgeIds.length === 1 ? 'Connection deleted' : `${edgeIds.length} connections deleted`)
|
|
}, [setEdges])
|
|
|
|
const handleDeleteSelectedEdges = useCallback(() => {
|
|
deleteEdgesById(selectedEdgeIds)
|
|
}, [deleteEdgesById, selectedEdgeIds])
|
|
|
|
const onEdgeContextMenu = useCallback((event: ReactMouseEvent, edge: Edge) => {
|
|
event.preventDefault()
|
|
event.stopPropagation()
|
|
deleteEdgesById([edge.id])
|
|
}, [deleteEdgesById])
|
|
|
|
const onEdgeDoubleClick = useCallback((event: ReactMouseEvent, edge: Edge) => {
|
|
event.preventDefault()
|
|
event.stopPropagation()
|
|
deleteEdgesById([edge.id])
|
|
}, [deleteEdgesById])
|
|
|
|
useEffect(() => {
|
|
const handleKeyDown = (event: KeyboardEvent) => {
|
|
const target = event.target as HTMLElement | null
|
|
const isEditingField =
|
|
target instanceof HTMLInputElement ||
|
|
target instanceof HTMLTextAreaElement ||
|
|
target instanceof HTMLSelectElement ||
|
|
Boolean(target?.closest('[contenteditable="true"]'))
|
|
|
|
if (isEditingField) return
|
|
if ((event.key === 'Delete' || event.key === 'Backspace') && selectedEdgeIds.length > 0) {
|
|
event.preventDefault()
|
|
deleteEdgesById(selectedEdgeIds)
|
|
}
|
|
}
|
|
|
|
window.addEventListener('keydown', handleKeyDown)
|
|
return () => window.removeEventListener('keydown', handleKeyDown)
|
|
}, [deleteEdgesById, selectedEdgeIds])
|
|
|
|
useEffect(() => {
|
|
if (selectedNodeId) {
|
|
setActiveUtilityTab('inspector')
|
|
return
|
|
}
|
|
if (activeUtilityTab === 'inspector') {
|
|
setActiveUtilityTab('library')
|
|
}
|
|
}, [activeUtilityTab, selectedNodeId])
|
|
|
|
const handleSave = useCallback(() => {
|
|
if (validation.errors.length > 0) {
|
|
toast.error('Resolve workflow validation errors before saving.')
|
|
return
|
|
}
|
|
onSave(buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode))
|
|
}, [edges, executionMode, nodes, onSave, validation.errors.length, workflow])
|
|
|
|
const handleDispatch = useCallback(() => {
|
|
if (!dispatchContextId.trim()) {
|
|
toast.error('Context ID is required for a graph test run.')
|
|
return
|
|
}
|
|
if (validation.errors.length > 0) {
|
|
toast.error('Resolve workflow validation errors before dispatching.')
|
|
return
|
|
}
|
|
dispatchMutation.mutate({
|
|
contextId: dispatchContextId.trim(),
|
|
config: buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode),
|
|
})
|
|
}, [dispatchContextId, dispatchMutation, edges, executionMode, nodes, validation.errors.length, workflow])
|
|
|
|
const handlePreflight = useCallback(() => {
|
|
if (!dispatchContextId.trim()) {
|
|
toast.error('Context ID is required for a graph preflight.')
|
|
return
|
|
}
|
|
if (validation.errors.length > 0) {
|
|
toast.error('Resolve workflow validation errors before running preflight.')
|
|
return
|
|
}
|
|
preflightMutation.mutate({
|
|
contextId: dispatchContextId.trim(),
|
|
config: buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode),
|
|
})
|
|
}, [dispatchContextId, edges, executionMode, nodes, preflightMutation, validation.errors.length, workflow])
|
|
|
|
const selectedNode = useMemo(
|
|
() => nodes.find(node => node.id === selectedNodeId),
|
|
[nodes, selectedNodeId],
|
|
)
|
|
|
|
return {
|
|
reactFlowWrapper,
|
|
nodeDefinitions,
|
|
nodeDefinitionsByStep,
|
|
nodes,
|
|
edges,
|
|
onNodesChange,
|
|
onEdgesChange,
|
|
selectedEdgeIds,
|
|
selectedNode,
|
|
workflowRuns,
|
|
selectedRun,
|
|
selectedRunComparison,
|
|
isComparisonLoading,
|
|
dispatchMutation,
|
|
preflightMutation,
|
|
dispatchContextId,
|
|
setDispatchContextId,
|
|
preflightResult,
|
|
executionMode,
|
|
setExecutionMode,
|
|
nodeMenuAnchor,
|
|
setNodeMenuAnchor,
|
|
activeUtilityTab,
|
|
setActiveUtilityTab,
|
|
validation,
|
|
graphFamily,
|
|
onConnect,
|
|
onNodeClick,
|
|
onEdgeClick,
|
|
onPaneClick,
|
|
handleParamsChange,
|
|
handlePipelineStepChange,
|
|
handlePaneContextMenu,
|
|
handleNodeContextMenu,
|
|
insertNode,
|
|
handleOpenToolbarNodeMenu,
|
|
handleAutoLayout,
|
|
handleDeleteSelectedEdges,
|
|
onEdgeContextMenu,
|
|
onEdgeDoubleClick,
|
|
handleSave,
|
|
handleDispatch,
|
|
handlePreflight,
|
|
setReactFlowInstance,
|
|
setSelectedRunId,
|
|
}
|
|
}
|