Update workflow sandbox mock then outputs generation

ee/codegen/src/__test__/__snapshots__/workflow-sandbox.test.ts.snap

@@ -1,5 +1,42 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
+exports[`Workflow Sandbox > write > should generate DatasetRow with mocks using array format for then_outputs 1`] = `
+"from vellum.workflows import MockNodeExecution
+from vellum.workflows.inputs import DatasetRow
+from vellum.workflows.sandbox import WorkflowSandboxRunner
+
+from .inputs import Inputs
+from .nodes.my_custom_node import MyCustomNode
+from .workflow import TestWorkflow
+
+dataset = [
+    DatasetRow(
+        label="Scenario with Array Mocks",
+        inputs=Inputs(test_input="test-value"),
+        mocks=[
+            MockNodeExecution(
+                when_condition=MyCustomNode.Execution.count.greater_than_or_equal_to(0),
+                then_outputs=MyCustomNode.Outputs(
+                    text_output="This is a mocked text output",
+                    json_output={
+                        "key": "mocked_value",
+                        "nested": {
+                            "data": 123,
+                        },
+                    },
+                ),
+            ),
+        ],
+    ),
+]
+
+runner = WorkflowSandboxRunner(workflow=TestWorkflow(), dataset=dataset)
+
+if __name__ == "__main__":
+    runner.run()
+"
+`;
+
 exports[`Workflow Sandbox > write > should generate DatasetRow with mocks when mocks are provided 1`] = `
 "from vellum.workflows import MockNodeExecution
 from vellum.workflows.inputs import DatasetRow

ee/codegen/src/__test__/workflow-sandbox.test.ts

@@ -430,5 +430,110 @@ describe("Workflow Sandbox", () => {
 
       expect(result).toMatchSnapshot();
     });
+
+    it("should generate DatasetRow with mocks using array format for then_outputs", async () => {
+      const writer = new Writer();
+      const uniqueWorkflowContext = workflowContextFactory();
+      const inputVariable: VellumVariable = {
+        id: "1",
+        key: "test_input",
+        type: "STRING",
+      };
+
+      uniqueWorkflowContext.addInputVariableContext(
+        inputVariableContextFactory({
+          inputVariableData: inputVariable,
+          workflowContext: uniqueWorkflowContext,
+        })
+      );
+
+      // Create a generic node with a known output id
+      const genericNodeData = genericNodeFactory({
+        nodeOutputs: [
+          {
+            id: "test-output-id-1",
+            name: "text_output",
+            type: "STRING",
+          },
+          {
+            id: "test-output-id-2",
+            name: "json_output",
+            type: "JSON",
+          },
+        ],
+      });
+      await nodeContextFactory({
+        workflowContext: uniqueWorkflowContext,
+        nodeData: genericNodeData,
+      });
+
+      const sandboxInputs: WorkflowSandboxDatasetRow[] = [
+        {
+          label: "Scenario with Array Mocks",
+          inputs: [
+            {
+              name: inputVariable.key,
+              type: "STRING",
+              value: "test-value",
+            },
+          ],
+          mocks: [
+            {
+              node_id: genericNodeData.id,
+              when_condition: {
+                type: "BINARY_EXPRESSION",
+                operator: ">=",
+                lhs: {
+                  type: "EXECUTION_COUNTER",
+                  nodeId: genericNodeData.id,
+                },
+                rhs: {
+                  type: "CONSTANT_VALUE",
+                  value: {
+                    type: "NUMBER",
+                    value: 0,
+                  },
+                },
+              },
+              then_outputs: [
+                {
+                  output_id: "test-output-id-1",
+                  value: {
+                    type: "CONSTANT_VALUE",
+                    value: {
+                      type: "STRING",
+                      value: "This is a mocked text output",
+                    },
+                  },
+                },
+                {
+                  output_id: "test-output-id-2",
+                  value: {
+                    type: "CONSTANT_VALUE",
+                    value: {
+                      type: "JSON",
+                      value: { key: "mocked_value", nested: { data: 123 } },
+                    },
+                  },
+                },
+              ],
+            },
+          ],
+        },
+      ];
+
+      const sandbox = codegen.workflowSandboxFile({
+        workflowContext: uniqueWorkflowContext,
+        sandboxInputs,
+      });
+
+      sandbox.write(writer);
+      const result = await writer.toStringFormatted();
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain("text_output=");
+      expect(result).toContain("json_output=");
+      expect(result).toContain("This is a mocked text output");
+    });
   });
 });

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

@@ -259,15 +259,47 @@ if __name__ == "__main__":
     }
 
     // Generate then_outputs by instantiating the node's Outputs class
-    const outputsArguments: MethodArgument[] = Object.entries(
-      mock.then_outputs
-    ).map(
-      ([key, value]) =>
-        new MethodArgument({
-          name: key,
-          value: new Json(value),
+    // Handle both array format (with output_id) and object format (with keys)
+    let outputsArguments: MethodArgument[];
+
+    if (Array.isArray(mock.then_outputs)) {
+      // Array format: each element has output_id and value
+      outputsArguments = mock.then_outputs
+        .map((output) => {
+          const outputName = nodeContext.getNodeOutputNameById(
+            output.output_id
+          );
+          if (isNil(outputName)) {
+            this.workflowContext.addError(
+              new NodeNotFoundError(
+                `Failed to find output name for output_id '${output.output_id}' in node '${mock.node_id}'`,
+                "WARNING"
+              )
+            );
+            return null;
+          }
+          return new MethodArgument({
+            name: outputName,
+            value: new WorkflowValueDescriptor({
+              workflowValueDescriptor: output.value,
+              workflowContext: this.workflowContext,
+            }),
+          });
         })
-    );
+        .filter((arg): arg is MethodArgument => !isNil(arg));
+    } else if (!isNil(mock.then_outputs)) {
+      // Object format: keys are output names, values are the output values
+      outputsArguments = Object.entries(mock.then_outputs).map(
+        ([key, value]) =>
+          new MethodArgument({
+            name: key,
+            value: new Json(value),
+          })
+      );
+    } else {
+      // No then_outputs provided
+      outputsArguments = [];
+    }
 
     const thenOutputsInstance = new ClassInstantiation({
       classReference: new Reference({

ee/codegen/src/types/vellum.ts

@@ -898,10 +898,15 @@ type WorkflowSandboxInput =
   | ImageInputRequest
   | DocumentInputRequest;
 export type WorkflowSandboxInputs = WorkflowSandboxInput[];
+export interface WorkflowSandboxDatasetRowMockOutput {
+  output_id: string;
+  value: WorkflowValueDescriptor;
+}
+
 export interface WorkflowSandboxDatasetRowMock {
   node_id: string;
   when_condition?: WorkflowValueDescriptor;
-  then_outputs: Record<string, unknown>;
+  then_outputs: Record<string, unknown> | WorkflowSandboxDatasetRowMockOutput[];
 }
 
 export type WorkflowSandboxDatasetRow =