Add disabled flag to MockNodeExecution (#3339)

* Add disabled flag to MockNodeExecution

- Add disabled: bool = False field to MockNodeExecution class
- Update runner to skip disabled mocks during evaluation
- Add disabled field to TypeScript types and serializers
- Update codegen to generate disabled parameter when true

Co-Authored-By: [email protected] <[email protected]>

* Fix serialization to exclude disabled=False from output

Co-Authored-By: [email protected] <[email protected]>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: [email protected] <[email protected]>

Author: devin-ai-integration[bot]

ee/codegen/src/generators/workflow-sandbox-file.ts

@@ -290,6 +290,16 @@ if __name__ == "__main__":
       })
     );
 
+    // Generate disabled parameter if it's true
+    if (mock.disabled === true) {
+      arguments_.push(
+        new MethodArgument({
+          name: "disabled",
+          value: python.TypeInstantiation.bool(true),
+        })
+      );
+    }
+
     return new ClassInstantiation({
       classReference: new Reference({
         name: "MockNodeExecution",

ee/codegen/src/serializers/vellum.ts

@@ -2601,6 +2601,7 @@ export declare namespace WorkflowSandboxDatasetRowMockSerializer {
     node_id: string;
     when_condition?: WorkflowValueDescriptorSerializer.Raw | null;
     then_outputs?: Record<string, unknown> | null;
+    disabled?: boolean | null;
   }
 }
 
@@ -2611,6 +2612,7 @@ export const WorkflowSandboxDatasetRowMockSerializer: ObjectSchema<
   node_id: stringSchema(),
   when_condition: WorkflowValueDescriptorSerializer.optional(),
   then_outputs: recordSchema(stringSchema(), unknownSchema()).optional(),
+  disabled: booleanSchema().optional(),
 });
 
 const WorkflowSandboxDatasetRowSerializer = undiscriminatedUnionSchema([

ee/codegen/src/types/vellum.ts

@@ -886,6 +886,7 @@ export interface WorkflowSandboxDatasetRowMock {
   node_id: string;
   when_condition?: WorkflowValueDescriptor;
   then_outputs?: Record<string, unknown>;
+  disabled?: boolean;
 }
 
 export type WorkflowSandboxDatasetRow =

src/vellum/workflows/nodes/mocks.py

@@ -103,6 +103,7 @@ class _RawMockWorkflowNodeConfig(UniversalBaseModel):
 class MockNodeExecution(UniversalBaseModel):
     when_condition: BaseDescriptor
     then_outputs: BaseOutputs
+    disabled: bool = False
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -112,6 +113,8 @@ def serialize_full_model(self, handler: Callable[[Any], Any], info: Serializatio
         serialized = handler(self)
         serialized["node_id"] = str(self.then_outputs.__class__.__parent_class__.__id__)
         serialized["type"] = "NODE_EXECUTION"
+        if not self.disabled:
+            del serialized["disabled"]
         return serialized
 
     @field_serializer("then_outputs")

src/vellum/workflows/runner/runner.py

@@ -510,6 +510,8 @@ def run_node(
             was_mocked: Optional[bool] = None
             mock_candidates = node_output_mocks_map.get(node.Outputs) or []
             for mock_candidate in mock_candidates:
+                if mock_candidate.disabled:
+                    continue
                 if mock_candidate.when_condition.resolve(node.state):
                     node_run_response = mock_candidate.then_outputs
                     was_mocked = True