feat(workflow): support dynamic input/output schemas in node definitions

Author: dingyi222666

packages/workflow/src/index.ts

@@ -71,7 +71,11 @@ function validateNodeInput(
 ): boolean {
     if (!definition.inputSchema) return true
     try {
-        definition.inputSchema.parse(input)
+        const schema =
+            typeof definition.inputSchema === 'function'
+                ? definition.inputSchema(node)
+                : definition.inputSchema
+        schema.parse(input)
         return true
     } catch {
         return false
@@ -81,13 +85,18 @@ function validateNodeInput(
 function validateNodeOutput(
     definition: NodeDefinition,
     output: NodeIO,
-    context: NodeContext
+    context: NodeContext,
+    node: WorkflowNode
 ): boolean {
     if (!definition.outputSchema) return true
     try {
-        for (const [key, schema] of Object.entries(definition.outputSchema)) {
+        const schema =
+            typeof definition.outputSchema === 'function'
+                ? definition.outputSchema(node)
+                : definition.outputSchema
+        for (const [key, schemaType] of Object.entries(schema)) {
             if (key in output) {
-                schema.parse(output[key])
+                schemaType.parse(output[key])
             }
         }
         return true
@@ -258,6 +267,11 @@ export async function executeWorkflow(
                     timing.duration = timing.endTime - timing.startTime
                     callbacks.onNodeError?.(node.id, node.type, error, timing)
                     failed.add(node.id)
+                    results[node.id] = {
+                        state: 'failed' as const,
+                        output: {},
+                        error
+                    }
                     return
                 }
 
@@ -272,7 +286,14 @@ export async function executeWorkflow(
                         timing.endTime = Date.now()
                         timing.duration = timing.endTime - timing.startTime
 
-                        if (!validateNodeOutput(definition, output, context)) {
+                        if (
+                            !validateNodeOutput(
+                                definition,
+                                output,
+                                context,
+                                node
+                            )
+                        ) {
                             throw new Error(
                                 `Invalid output for node ${node.id}`
                             )

packages/workflow/src/types.ts

@@ -30,8 +30,12 @@ export type NodeDefinition<
     TData = unknown
 > = {
     run: (input: TInput, context: NodeContext, data?: TData) => Promise<TOutput>
-    inputSchema?: z.ZodType<TInput>
-    outputSchema?: Record<string, z.ZodType<unknown>>
+    inputSchema?:
+        | z.ZodType<TInput>
+        | ((node: WorkflowNode) => z.ZodType<TInput>)
+    outputSchema?:
+        | Record<string, z.ZodType<unknown>>
+        | ((node: WorkflowNode) => Record<string, z.ZodType<unknown>>)
     dataSchema?: z.ZodType<TData>
 }
 

packages/workflow/tests/example.spec.ts

@@ -143,13 +143,17 @@ describe('Workflow with Optional Outputs and Skipped Nodes', () => {
         }
 
         const [lastNodeResult, allResults] = await runExample()
-        
+
         // Verify last node result (evenHandler1)
         expect(lastNodeResult.state).to.equal('completed')
         expect(lastNodeResult.output).to.deep.equal({ doubled: 8 })
-        
+
         // Verify all node results
-        expect(allResults).to.have.all.keys('processor1', 'evenHandler1', 'oddHandler1')
+        expect(allResults).to.have.all.keys(
+            'processor1',
+            'evenHandler1',
+            'oddHandler1'
+        )
         expect(allResults.processor1.state).to.equal('completed')
         expect(allResults.processor1.output).to.deep.equal({ even: 4 })
         expect(allResults.evenHandler1.state).to.equal('completed')
@@ -284,7 +288,11 @@ describe('Workflow with Optional Outputs and Skipped Nodes', () => {
         expect(lastNodeResult.output).to.deep.equal({ doubled: 10 })
 
         // Verify all node results
-        expect(allResults).to.have.all.keys('if1', 'trueHandler1', 'falseHandler1')
+        expect(allResults).to.have.all.keys(
+            'if1',
+            'trueHandler1',
+            'falseHandler1'
+        )
         expect(allResults.if1.state).to.equal('completed')
         expect(allResults.if1.output).to.deep.equal({ true: 5 })
         expect(allResults.trueHandler1.state).to.equal('completed')
@@ -396,4 +404,98 @@ describe('Workflow with Optional Outputs and Skipped Nodes', () => {
         expect(allResults.divide.state).to.equal('completed')
         expect(allResults.divide.output).to.deep.equal({ result: 5 })
     })
+
+    it('should handle dynamic input/output schemas based on node config', async () => {
+        // Define node types with dynamic schemas
+        type EvalInput = { value: number | string }
+        type EvalOutput = { result: unknown }
+
+        const evalNode: NodeDefinition<EvalInput, EvalOutput> = {
+            inputSchema: (node: WorkflowNode) => {
+                // Dynamic input schema based on node config
+                const inputType = node.config?.inputType || 'number'
+                return z.object({
+                    value: inputType === 'string' ? z.string() : z.number()
+                })
+            },
+            outputSchema: (node: WorkflowNode) => {
+                // Dynamic output schema based on node config
+                const outputType = node.config?.outputType || 'number'
+                return {
+                    result: outputType === 'string' ? z.string() : z.number()
+                }
+            },
+            run: async (input: EvalInput) => {
+                // Simple evaluation for demonstration
+                const result = input.value.toString()
+                return { result }
+            }
+        }
+
+        // Create and configure workflow
+        const factory = createNodeFactory()
+        factory.registerNode('eval', evalNode)
+
+        const workflow: WorkflowNode[] = [
+            {
+                id: 'eval1',
+                type: 'eval',
+                dependencies: [],
+                config: {
+                    inputType: 'number',
+                    outputType: 'string'
+                }
+            }
+        ]
+
+        const initialContext: NodeContext = {
+            variables: {
+                value: 42
+            },
+            metadata: {}
+        }
+
+        const [lastNodeResult, allResults] = await executeWorkflow(
+            workflow,
+            factory,
+            initialContext,
+            {
+                maxRetries: 2,
+                maxParallel: 2
+            }
+        )
+
+        // Verify results
+        expect(lastNodeResult.state).to.equal('completed')
+        expect(lastNodeResult.output).to.deep.equal({ result: '42' })
+        expect(allResults.eval1.state).to.equal('completed')
+        expect(allResults.eval1.output).to.deep.equal({ result: '42' })
+
+        // Test with invalid input type
+        const invalidWorkflow: WorkflowNode[] = [
+            {
+                id: 'eval2',
+                type: 'eval',
+                dependencies: [],
+                config: {
+                    inputType: 'string',
+                    outputType: 'number'
+                }
+            }
+        ]
+
+        const [invalidLastResult, invalidResults] = await executeWorkflow(
+            invalidWorkflow,
+            factory,
+            initialContext,
+            {
+                maxRetries: 2,
+                maxParallel: 2
+            }
+        )
+
+        // Verify validation failure
+        expect(invalidLastResult.state).to.equal('failed')
+        expect(invalidResults.eval2.state).to.equal('failed')
+    })
 })