[APO-2121] Codegen workflow trigger id for scenarios (#3067)

* [APO-2121] Codegen workflow trigger id for scenarios

* update to generate trigger class instead of id

* npm run lint

---------

Co-authored-by: Harrison Ngo <[email protected]>

Author: NgoHarrison

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

@@ -1,5 +1,31 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
+exports[`Workflow Sandbox > write > should generate DatasetRow with trigger when workflow_trigger_id is provided 1`] = `
+"from vellum.workflows.inputs import DatasetRow
+from vellum.workflows.sandbox import WorkflowSandboxRunner
+
+from .inputs import Inputs
+from .triggers.scheduled import ScheduleTrigger
+from .workflow import TestWorkflow
+
+dataset = [
+    DatasetRow(
+        label="Scenario with Trigger ID",
+        inputs=Inputs(test_input="test-value"),
+        trigger=ScheduleTrigger,
+    ),
+    DatasetRow(
+        label="Scenario without Trigger ID", inputs=Inputs(test_input="test-value-2")
+    ),
+]
+
+runner = WorkflowSandboxRunner(workflow=TestWorkflow(), dataset=dataset)
+
+if __name__ == "__main__":
+    runner.run()
+"
+`;
+
 exports[`Workflow Sandbox > write > should generate DatasetRow without inputs parameter when dataset row has label but no inputs 1`] = `
 "from vellum.workflows.inputs import DatasetRow
 from vellum.workflows.sandbox import WorkflowSandboxRunner

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

@@ -6,7 +6,11 @@ import { workflowContextFactory } from "./helpers";
 import { inputVariableContextFactory } from "./helpers/input-variable-context-factory";
 
 import * as codegen from "src/codegen";
-import { WorkflowSandboxDatasetRow } from "src/types/vellum";
+import {
+  WorkflowSandboxDatasetRow,
+  WorkflowTrigger,
+  WorkflowTriggerType,
+} from "src/types/vellum";
 
 describe("Workflow Sandbox", () => {
   const generateSandboxFile = async (
@@ -172,5 +176,74 @@ describe("Workflow Sandbox", () => {
       expect(result).toContain('DatasetRow(label="Test Label Only")');
       expect(result).not.toContain("inputs=");
     });
+
+    it("should generate DatasetRow with trigger when workflow_trigger_id is provided", async () => {
+      const writer = new Writer();
+      const triggerId = "550e8400-e29b-41d4-a716-446655440000";
+      const triggers: WorkflowTrigger[] = [
+        {
+          id: triggerId,
+          type: WorkflowTriggerType.SCHEDULED,
+          attributes: [],
+          cron: "* * * * *",
+          timezone: "UTC",
+        },
+      ];
+      const uniqueWorkflowContext = workflowContextFactory({ triggers });
+      const inputVariable: VellumVariable = {
+        id: "1",
+        key: "test_input",
+        type: "STRING",
+      };
+
+      uniqueWorkflowContext.addInputVariableContext(
+        inputVariableContextFactory({
+          inputVariableData: inputVariable,
+          workflowContext: uniqueWorkflowContext,
+        })
+      );
+
+      const sandboxInputs: WorkflowSandboxDatasetRow[] = [
+        {
+          label: "Scenario with Trigger ID",
+          inputs: [
+            {
+              name: inputVariable.key,
+              type: "STRING",
+              value: "test-value",
+            },
+          ],
+          workflow_trigger_id: triggerId,
+        },
+        {
+          label: "Scenario without Trigger ID",
+          inputs: [
+            {
+              name: inputVariable.key,
+              type: "STRING",
+              value: "test-value-2",
+            },
+          ],
+        },
+      ];
+
+      const sandbox = codegen.workflowSandboxFile({
+        workflowContext: uniqueWorkflowContext,
+        sandboxInputs,
+      });
+
+      sandbox.write(writer);
+      const result = await writer.toStringFormatted();
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain("trigger=ScheduleTrigger");
+      const lines = result.split("\n");
+      const secondDatasetRowIndex = lines.findIndex((line) =>
+        line.includes('label="Scenario without Trigger ID"')
+      );
+      expect(secondDatasetRowIndex).toBeGreaterThan(-1);
+      const secondDatasetRowLine = lines[secondDatasetRowIndex];
+      expect(secondDatasetRowLine).not.toContain("trigger=");
+    });
   });
 });

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

@@ -10,6 +10,7 @@ import {
 } from "src/types/vellum";
 import { removeEscapeCharacters } from "src/utils/casing";
 import { getGeneratedInputsModulePath } from "src/utils/paths";
+import { getTriggerClassInfo } from "src/utils/triggers";
 
 export declare namespace WorkflowSandboxFile {
   interface Args extends BasePersistedFile.Args {
@@ -98,6 +99,9 @@ if __name__ == "__main__":
     const label: string = Array.isArray(row)
       ? `Example ${index + 1}`
       : row.label;
+    const workflowTriggerId: string | undefined = Array.isArray(row)
+      ? undefined
+      : row.workflow_trigger_id;
 
     const hasInputs = inputs.length > 0;
     const arguments_: python.MethodArgument[] = [
@@ -147,6 +151,19 @@ if __name__ == "__main__":
       );
     }
 
+    if (!isNil(workflowTriggerId)) {
+      const triggerReference = this.getTriggerReference(workflowTriggerId);
+
+      if (triggerReference) {
+        arguments_.push(
+          python.methodArgument({
+            name: "trigger",
+            value: triggerReference,
+          })
+        );
+      }
+    }
+
     return python.instantiateClass({
       classReference: python.reference({
         name: "DatasetRow",
@@ -155,4 +172,32 @@ if __name__ == "__main__":
       arguments_,
     });
   }
+
+  private getTriggerReference(
+    workflowTriggerId: string
+  ): python.Reference | undefined {
+    const triggerContext =
+      this.workflowContext.findTriggerContext(workflowTriggerId);
+
+    if (triggerContext) {
+      return python.reference({
+        name: triggerContext.triggerClassName,
+        modulePath: triggerContext.triggerModulePath,
+      });
+    }
+
+    const triggers = this.workflowContext.triggers ?? [];
+    const trigger = triggers.find((t) => t.id === workflowTriggerId);
+
+    if (!trigger) {
+      return undefined;
+    }
+
+    const triggerClassInfo = getTriggerClassInfo(trigger, this.workflowContext);
+
+    return python.reference({
+      name: triggerClassInfo.className,
+      modulePath: triggerClassInfo.modulePath,
+    });
+  }
 }

ee/codegen/src/types/vellum.ts

@@ -888,8 +888,12 @@ type WorkflowSandboxInput =
 export type WorkflowSandboxInputs = WorkflowSandboxInput[];
 export type WorkflowSandboxDatasetRow =
   | WorkflowSandboxInputs
-  | { label: string; inputs: WorkflowSandboxInputs }
-  | { label: string };
+  | {
+      label: string;
+      inputs: WorkflowSandboxInputs;
+      workflow_trigger_id?: string;
+    }
+  | { label: string; workflow_trigger_id?: string };
 
 export interface UnaryWorkflowExpression {
   type: "UNARY_EXPRESSION";