refactor(workflow): move output validation off inline node paths

Make the scheduler the single authoritative point for output
validation. FunctionNode.wrappedFn and ToolNode.runTool no longer
validate the output against their schema inline; the schema still
flows through BaseNode so the scheduler enforces it on every yielded
event with non-nil output. This removes duplicate validation and the
risk of drift as the scheduler logic evolves (symmetric to the input
validation cleanup).

The FunctionTool {"result": X} unwrap fallback is not lost: it moves
to a ToolNode.ValidateOutput override, the correct home for that
tool-specific convention. The override tries standard validation
first, then unwraps a "result" key and validates the unwrapped value,
and finally re-runs standard validation so callers see the original
schema-mismatch error rather than a fallback artifact. nil schema is
a passthrough.

Tests that previously asserted on inline behavior are updated:
- FunctionNode: the inline ValidationError table case is replaced by
  TestFunctionNode_ValidateOutput, which asserts Run passes the output
  through unchanged and ValidateOutput surfaces the schema mismatch.
- ToolNode: TestToolNode_Run now asserts the raw FunctionTool map
  output from Run (no inline unwrap); the new
  TestToolNode_ValidateOutput covers the four override cases
  (direct-valid passthrough, {"result": X} unwrap, both-fail original
  error, nil-schema passthrough).

BUG=516382303
BUG=516382092

Author: wolo-lab

workflow/function_node.go

@@ -230,11 +230,6 @@ func NewFunctionNodeFromState[Params, OUT any](
 		if err != nil {
 			return output, err
 		}
-		if oschema != nil {
-			if err := typeutil.ValidateWithJSONSchema(output, oschema); err != nil {
-				return nil, fmt.Errorf("function node %s: validation failed for output %T: %w", name, new(OUT), err)
-			}
-		}
 		return output, nil
 	}
 
@@ -276,13 +271,6 @@ func newFunctionNodeWithResolvedSchemas[IN, OUT any](name string, fn func(ctx ag
 			return output, err
 		}
 
-		if outputSchema != nil {
-			validateErr := typeutil.ValidateWithJSONSchema(output, outputSchema)
-			if validateErr != nil {
-				return nil, fmt.Errorf("function node %s: validation failed for output %T: %w", name, new(OUT), validateErr)
-			}
-		}
-
 		return output, nil
 	}
 
@@ -321,11 +309,6 @@ func newEmittingFunctionNodeWithResolvedSchemas[IN, OUT any](name string, fn Emi
 		if err != nil {
 			return nil, err
 		}
-		if outputSchema != nil {
-			if validateErr := typeutil.ValidateWithJSONSchema(output, outputSchema); validateErr != nil {
-				return nil, fmt.Errorf("function node %s: validation failed for output %T: %w", name, new(OUT), validateErr)
-			}
-		}
 		return output, nil
 	}
 

workflow/function_node_test.go

@@ -37,9 +37,6 @@ func TestNewFunctionNodeWithSchema(t *testing.T) {
 	type Output struct {
 		Result string `json:"result"`
 	}
-	type TargetOutput struct {
-		Result int `json:"result"`
-	}
 
 	tests := []struct {
 		name         string
@@ -79,18 +76,6 @@ func TestNewFunctionNodeWithSchema(t *testing.T) {
 			wantOutput:   map[string]any{"result": "zero"},
 			wantErr:      false,
 		},
-		{
-			name:     "ValidationError",
-			nodeName: "test",
-			fn: func(ctx agent.Context, input Input) (map[string]any, error) {
-				return map[string]any{"result": "not-an-int"}, nil
-			},
-			inputSchema:  mustSchema[Input](t),
-			outputSchema: mustSchema[TargetOutput](t),
-			input:        Input{Value: "hello"},
-			wantErr:      true,
-			errSubstr:    "validation failed for output",
-		},
 	}
 
 	for _, tc := range tests {
@@ -132,6 +117,49 @@ func TestNewFunctionNodeWithSchema(t *testing.T) {
 	}
 }
 
+// TestFunctionNode_ValidateOutput verifies that output schema validation
+// is enforced through the node-level ValidateOutput contract (invoked
+// scheduler-side), not inline inside Run. Run itself passes the output
+// through unchanged; ValidateOutput surfaces the schema mismatch.
+func TestFunctionNode_ValidateOutput(t *testing.T) {
+	type Input struct {
+		Value string `json:"value"`
+	}
+	type TargetOutput struct {
+		Result int `json:"result"`
+	}
+
+	fn := func(ctx agent.Context, input Input) (map[string]any, error) {
+		return map[string]any{"result": "not-an-int"}, nil
+	}
+	node, err := NewFunctionNodeWithSchema[Input, map[string]any](
+		"test", fn, mustSchema[Input](t), mustSchema[TargetOutput](t), defaultNodeConfig)
+	if err != nil {
+		t.Fatalf("NewFunctionNodeWithSchema failed: %v", err)
+	}
+
+	// Run no longer validates: it yields the raw output without error.
+	mockCtx := &MockInvocationContext{sess: nil}
+	exCtx := agent.NewNodeContext(mockCtx, nil)
+	var got any
+	count := 0
+	for ev, err := range node.Run(exCtx, Input{Value: "hello"}) {
+		if err != nil {
+			t.Fatalf("Run returned unexpected error: %v", err)
+		}
+		got = ev.Output
+		count++
+	}
+	if count != 1 {
+		t.Fatalf("expected 1 event from Run, got %d", count)
+	}
+
+	// ValidateOutput (the scheduler-side contract) rejects the mismatch.
+	if _, err := node.ValidateOutput(got); err == nil {
+		t.Fatalf("ValidateOutput: expected validation error, got nil")
+	}
+}
+
 func mustSchema[T any](t *testing.T) *jsonschema.Schema {
 	t.Helper()
 	s, err := jsonschema.For[T](nil)

workflow/tool_node.go

@@ -112,23 +112,38 @@ func (n *ToolNode) runTool(toolCtx agent.Context, input any) (any, error) {
 		return nil, fmt.Errorf("tool %q execution failed: %w", n.tool.Name(), err)
 	}
 
-	var toolOutput any = output
-
-	// Validate
-	if schema := n.OutputSchema(); schema != nil {
-		if err := schema.Validate(output); err != nil {
-			if val, ok := output["result"]; ok {
-				if err := schema.Validate(val); err != nil {
-					return nil, fmt.Errorf("converting tool %q output: validation failed for result key: %w", n.tool.Name(), err)
-				}
-				toolOutput = val
-			} else {
-				return nil, fmt.Errorf("converting tool %q output: validation failed: %w", n.tool.Name(), err)
+	return output, nil
+}
+
+// ValidateOutput validates the tool output against the node's output
+// schema, adding a FunctionTool-specific fallback on top of the default
+// behavior: when the output is a map of shape {"result": X} that fails
+// direct schema validation, it retries against the unwrapped "result"
+// value and, on success, returns that unwrapped value.
+//
+// This override is the home for the {"result": X} convention because it
+// is tool-specific; making it a general default could mask genuine
+// validation errors in other node types.
+func (n *ToolNode) ValidateOutput(out any) (any, error) {
+	schema := n.OutputSchema()
+	if schema == nil {
+		return out, nil
+	}
+	// Try standard validation first.
+	if validated, err := defaultValidateOutput(out, schema); err == nil {
+		return validated, nil
+	}
+	// Fallback: unwrap {"result": X} (FunctionTool convention).
+	if m, ok := out.(map[string]any); ok {
+		if val, ok := m["result"]; ok {
+			if validated, err := defaultValidateOutput(val, schema); err == nil {
+				return validated, nil
 			}
 		}
 	}
-
-	return toolOutput, nil
+	// Re-run standard validation so the caller sees the original
+	// schema-mismatch error rather than a fallback artifact.
+	return defaultValidateOutput(out, schema)
 }
 
 // Run implements the Node interface and executes the tool.

workflow/tool_node_test.go

@@ -110,9 +110,6 @@ func TestToolNode_Run(t *testing.T) {
 	type Output struct {
 		Greeting string `json:"greeting"`
 	}
-	type ErrorOutput struct {
-		Result int `json:"result"`
-	}
 
 	tests := []struct {
 		name      string
@@ -162,26 +159,14 @@ func TestToolNode_Run(t *testing.T) {
 			node: func(t tool.Tool) (Node, error) {
 				return NewToolNodeTyped[Input, string](t, defaultNodeConfig)
 			},
+			// Run no longer unwraps {"result": X}; that happens
+			// scheduler-side in ToolNode.ValidateOutput. Run yields the
+			// raw FunctionTool map output.
 			extract: func(t *testing.T, out any) string {
-				return out.(string)
+				return out.(map[string]any)["result"].(string)
 			},
 			want: "HELLO WORLD",
 		},
-		{
-			name: "schema_validation_error",
-			tool: func() (tool.Tool, error) {
-				return functiontool.New(functiontool.Config{
-					Name: "test_tool",
-				}, func(ctx agent.Context, in map[string]any) (map[string]any, error) {
-					return map[string]any{"result": "not-an-int"}, nil
-				})
-			},
-			nodeInput: map[string]any{},
-			node: func(t tool.Tool) (Node, error) {
-				return NewToolNodeTyped[map[string]any, ErrorOutput](t, defaultNodeConfig)
-			},
-			wantErr: "converting tool \"test_tool\" output",
-		},
 		{
 			name: "tool_execution_error",
 			tool: func() (tool.Tool, error) {
@@ -276,6 +261,78 @@ func TestToolNode_Run(t *testing.T) {
 	}
 }
 
+// TestToolNode_ValidateOutput exercises the FunctionTool-specific
+// {"result": X} unwrap fallback that ToolNode layers on top of the
+// default schema validation.
+func TestToolNode_ValidateOutput(t *testing.T) {
+	type Result struct {
+		Greeting string `json:"greeting"`
+	}
+
+	// Node carrying a Result output schema.
+	schemaNode := &ToolNode{
+		BaseNode: NewBaseNodeWithSchemas(
+			"greet", "", defaultNodeConfig, nil, resolveTestSchema[Result](t)),
+	}
+	// Node with no output schema.
+	nilSchemaNode := &ToolNode{
+		BaseNode: NewBaseNode("greet", "", defaultNodeConfig),
+	}
+
+	valid := map[string]any{"greeting": "Hello World"}
+
+	tests := []struct {
+		name    string
+		node    *ToolNode
+		output  any
+		want    any
+		wantErr bool
+	}{
+		{
+			name:   "direct_valid_passes_through",
+			node:   schemaNode,
+			output: valid,
+			want:   valid,
+		},
+		{
+			name:   "result_wrapped_is_unwrapped",
+			node:   schemaNode,
+			output: map[string]any{"result": valid},
+			want:   valid,
+		},
+		{
+			name:    "fails_direct_and_fallback",
+			node:    schemaNode,
+			output:  map[string]any{"result": map[string]any{"unexpected": 1}},
+			wantErr: true,
+		},
+		{
+			name:   "nil_schema_passes_through",
+			node:   nilSchemaNode,
+			output: map[string]any{"anything": 1},
+			want:   map[string]any{"anything": 1},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := tc.node.ValidateOutput(tc.output)
+			if tc.wantErr {
+				if err == nil {
+					t.Fatalf("ValidateOutput: expected error, got nil (out=%v)", got)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("ValidateOutput: unexpected error: %v", err)
+			}
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("ValidateOutput mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
 func TestToolNode_WorkflowIntegration(t *testing.T) {
 	type Input struct {
 		Val int `json:"val"`