feat(workflow): enhance node definition and execution with data handling and validation

Author: dingyi222666

packages/workflow/src/index.ts

@@ -7,8 +7,10 @@ import {
     NodeId,
     NodeIO,
     NodeResult,
+    WorkflowCallbacks,
     WorkflowNode,
-    WorkflowOptions
+    WorkflowOptions,
+    WorkflowResult
 } from './types.ts'
 
 export * from './types.ts'
@@ -17,12 +19,17 @@ export function createNodeFactory(): NodeFactory {
     const nodes = new Map<string, NodeDefinition>()
 
     return {
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        registerNode: (type: string, definition: NodeDefinition<any, any>) => {
+        registerNode: (
+            type: string,
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            definition: NodeDefinition<any, any, any>
+        ) => {
             nodes.set(type, definition)
         },
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        getNode: (type: string) => nodes.get(type) as NodeDefinition<any, any>
+
+        getNode: (type: string) =>
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            nodes.get(type) as NodeDefinition<any, any, any>
     }
 }
 
@@ -56,28 +63,116 @@ function detectCircularDependencies(nodes: WorkflowNode[]): boolean {
     return nodes.some((node) => hasCycle(node.id))
 }
 
+function validateNodeInput(
+    node: WorkflowNode,
+    definition: NodeDefinition,
+    input: NodeIO,
+    context: NodeContext
+): boolean {
+    if (!definition.inputSchema) return true
+    try {
+        definition.inputSchema.parse(input)
+        return true
+    } catch {
+        return false
+    }
+}
+
+function validateNodeOutput(
+    definition: NodeDefinition,
+    output: NodeIO,
+    context: NodeContext
+): boolean {
+    if (!definition.outputSchema) return true
+    try {
+        for (const [key, schema] of Object.entries(definition.outputSchema)) {
+            if (key in output) {
+                schema.parse(output[key])
+            }
+        }
+        return true
+    } catch {
+        return false
+    }
+}
+
+function checkAndMarkSkippedNodes(
+    nodes: WorkflowNode[],
+    completed: Set<NodeId>,
+    failed: Set<NodeId>,
+    skipped: Set<NodeId>,
+    results: Record<NodeId, NodeResult>,
+    callbacks?: WorkflowCallbacks
+): void {
+    nodes.forEach((node) => {
+        if (
+            completed.has(node.id) ||
+            failed.has(node.id) ||
+            skipped.has(node.id)
+        ) {
+            return
+        }
+
+        const cannotBeExecuted = node.dependencies.some((dep) => {
+            const depId = resolveDependency(dep)
+            if (failed.has(depId)) return true
+
+            if (completed.has(depId) && typeof dep === 'object' && dep.portId) {
+                const depOutput = results[depId]?.output || {}
+                return depOutput[dep.portId] === undefined
+            }
+
+            return false
+        })
+
+        if (cannotBeExecuted) {
+            skipped.add(node.id)
+            results[node.id] = {
+                state: 'skipped' as const,
+                output: {}
+            }
+            callbacks?.onNodeSkipped?.(node.id, node.type)
+        }
+    })
+}
+
 function getExecutableNodes(
     nodes: WorkflowNode[],
     completed: Set<NodeId>,
-    failed: Set<NodeId>
+    failed: Set<NodeId>,
+    skipped: Set<NodeId>,
+    results: Record<NodeId, NodeResult>
 ): WorkflowNode[] {
     return nodes.filter((node) => {
-        if (completed.has(node.id) || failed.has(node.id)) return false
+        if (
+            completed.has(node.id) ||
+            failed.has(node.id) ||
+            skipped.has(node.id)
+        )
+            return false
+
         return node.dependencies.every((dep) => {
             const depId = resolveDependency(dep)
-            return completed.has(depId) && !failed.has(depId)
+            const depResult = completed.has(depId)
+
+            if (typeof dep === 'object' && dep.portId) {
+                const depOutput = results[depId]?.output || {}
+                return depResult && depOutput[dep.portId] !== undefined
+            }
+
+            return depResult && !failed.has(depId)
         })
     })
 }
 
-function validateNodeInput(
+function validateNodeData(
     node: WorkflowNode,
     definition: NodeDefinition,
-    input: NodeIO
+    context: NodeContext
 ): boolean {
-    if (!definition.inputSchema) return true
+    if (!definition.dataSchema || !node.data) return true
     try {
-        definition.inputSchema.parse(input)
+        definition.dataSchema.parse(node.data)
         return true
     } catch {
         return false
@@ -89,25 +184,34 @@ export async function executeWorkflow(
     factory: NodeFactory,
     initialContext: NodeContext = { variables: {}, metadata: {} },
     options: WorkflowOptions = {}
-): Promise<Record<NodeId, NodeResult>> {
+): Promise<WorkflowResult> {
     const { maxRetries = 3, maxParallel = 4, callbacks = {} } = options
+    const results: Record<NodeId, NodeResult> = {}
 
     if (detectCircularDependencies(nodes)) {
         throw new Error('Circular dependencies detected in workflow')
     }
 
     const completed = new Set<NodeId>()
     const failed = new Set<NodeId>()
-    const results: Record<NodeId, NodeResult> = {}
+    const skipped = new Set<NodeId>()
     const context: NodeContext = { ...initialContext }
+    let lastExecutedNode: WorkflowNode | undefined
 
-    while (completed.size + failed.size < nodes.length) {
-        const executableNodes = getExecutableNodes(nodes, completed, failed)
+    while (completed.size + failed.size + skipped.size < nodes.length) {
+        const executableNodes = getExecutableNodes(
+            nodes,
+            completed,
+            failed,
+            skipped,
+            results
+        )
         if (executableNodes.length === 0) break
 
         const executions = executableNodes
             .slice(0, maxParallel)
             .map(async (node) => {
+                lastExecutedNode = node
                 const definition = factory.getNode(node.type)
                 if (!definition) {
                     failed.add(node.id)
@@ -120,6 +224,15 @@ export async function executeWorkflow(
 
                 callbacks.onNodeStart?.(node.id, node.type)
 
+                if (!validateNodeData(node, definition, context)) {
+                    const error = new Error(`Invalid data for node ${node.id}`)
+                    timing.endTime = Date.now()
+                    timing.duration = timing.endTime - timing.startTime
+                    callbacks.onNodeError?.(node.id, node.type, error, timing)
+                    failed.add(node.id)
+                    return
+                }
+
                 const input =
                     node.dependencies.length === 0
                         ? { ...context.variables }
@@ -138,7 +251,8 @@ export async function executeWorkflow(
                               return { ...acc, ...depOutput }
                           }, {} as NodeIO)
 
-                if (!validateNodeInput(node, definition, input)) {
+                if (!validateNodeInput(node, definition, input, context)) {
+                    console.log(node, input)
                     const error = new Error(`Invalid input for node ${node.id}`)
                     timing.endTime = Date.now()
                     timing.duration = timing.endTime - timing.startTime
@@ -150,12 +264,22 @@ export async function executeWorkflow(
                 let retries = 0
                 while (retries <= maxRetries) {
                     try {
-                        const output = await definition.run(input, context)
+                        const output = await definition.run(
+                            input,
+                            context,
+                            node.data
+                        )
                         timing.endTime = Date.now()
                         timing.duration = timing.endTime - timing.startTime
 
+                        if (!validateNodeOutput(definition, output, context)) {
+                            throw new Error(
+                                `Invalid output for node ${node.id}`
+                            )
+                        }
+
                         const result: NodeResult = {
-                            state: 'completed',
+                            state: 'completed' as const,
                             output
                         }
 
@@ -167,6 +291,16 @@ export async function executeWorkflow(
                             result,
                             timing
                         )
+
+                        // Check for nodes that should be skipped after this node completes
+                        checkAndMarkSkippedNodes(
+                            nodes,
+                            completed,
+                            failed,
+                            skipped,
+                            results,
+                            callbacks
+                        )
                         break
                     } catch (error) {
                         retries++
@@ -175,7 +309,7 @@ export async function executeWorkflow(
                             timing.duration = timing.endTime - timing.startTime
 
                             const result: NodeResult = {
-                                state: 'failed',
+                                state: 'failed' as const,
                                 output: {},
                                 error: error as Error
                             }
@@ -195,6 +329,13 @@ export async function executeWorkflow(
         await Promise.all(executions)
     }
 
+    const lastNodeResult: NodeResult = lastExecutedNode
+        ? results[lastExecutedNode.id]
+        : {
+              state: 'skipped' as const,
+              output: {}
+          }
+
     callbacks.onWorkflowComplete?.(results)
-    return results
+    return [lastNodeResult, results]
 }

packages/workflow/src/types.ts

@@ -22,13 +22,17 @@ export type NodeResult = {
     error?: Error
 }
 
+export type WorkflowResult = [NodeResult, Record<NodeId, NodeResult>]
+
 export type NodeDefinition<
     TInput extends NodeIO = NodeIO,
-    TOutput extends NodeIO = NodeIO
+    TOutput extends NodeIO = NodeIO,
+    TData = unknown
 > = {
-    run: (input: TInput, context: NodeContext) => Promise<TOutput>
+    run: (input: TInput, context: NodeContext, data?: TData) => Promise<TOutput>
     inputSchema?: z.ZodType<TInput>
-    outputSchema?: z.ZodType<TOutput>
+    outputSchema?: Record<string, z.ZodType<unknown>>
+    dataSchema?: z.ZodType<TData>
 }
 
 export type NodeDependency =
@@ -44,6 +48,7 @@ export type WorkflowNode = {
     type: string
     dependencies: NodeDependency[]
     config?: Record<string, unknown>
+    data?: unknown
 }
 
 export type NodeExecutionTiming = {
@@ -82,15 +87,17 @@ export type NodeFactory = {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     registerNode<
         TInput extends NodeIO = NodeIO<unknown>,
-        TOutput extends NodeIO = NodeIO<unknown>
+        TOutput extends NodeIO = NodeIO<unknown>,
+        TData = unknown
     >(
         type: string,
-        definition: NodeDefinition<TInput, TOutput>
+        definition: NodeDefinition<TInput, TOutput, TData>
     ): void
     getNode<
         TInput extends NodeIO = NodeIO<unknown>,
-        TOutput extends NodeIO = NodeIO<unknown>
+        TOutput extends NodeIO = NodeIO<unknown>,
+        TData = unknown
     >(
         type: string
-    ): NodeDefinition<TInput, TOutput> | undefined
+    ): NodeDefinition<TInput, TOutput, TData> | undefined
 }