[APO-2825] Add codegen test for state variables in inline subworkflow nodes

ee/codegen/src/__test__/__snapshots__/generate-code.test.ts.snap

@@ -153,6 +153,28 @@ class WorkflowDisplay(BaseWorkflowDisplay[Workflow]):
 "
 `;
 
+exports[`generateCode > should generate code for %1 inline-subworkflow-with-state-variables.ts > nodes/subworkflow_with_state/state.py 1`] = `
+"from typing import Optional, Union
+
+from vellum.workflows.state import BaseState
+
+
+class State(BaseState):
+    inner_counter: Optional[Union[float, int]] = 0
+"
+`;
+
+exports[`generateCode > should generate code for %1 inline-subworkflow-with-state-variables.ts > state.py 1`] = `
+"from typing import Optional, Union
+
+from vellum.workflows.state import BaseState
+
+
+class State(BaseState):
+    parent_counter: Optional[Union[float, int]] = 0
+"
+`;
+
 exports[`generateCode > should generate code for %1 integration-trigger-with-entrypoint-and-outgoing-edges.ts > triggers/slack_message_trigger.py 1`] = `
 "from vellum.workflows.triggers import IntegrationTrigger
 

ee/codegen/src/__test__/generate-code-fixtures/inline-subworkflow-with-state-variables.ts

@@ -0,0 +1,314 @@
+import { v4 as uuidv4 } from "uuid";
+
+const outputVariableId = uuidv4();
+const outputNodeId = uuidv4();
+const outputInputId = uuidv4();
+const outputHandleId = uuidv4();
+const entrypointNodeId = uuidv4();
+const entrypointHandleId = uuidv4();
+
+const parentStateVariableId = uuidv4();
+const innerStateVariableId = uuidv4();
+
+const subworkflowNodeId = uuidv4();
+const subworkflowSourceHandleId = uuidv4();
+const subworkflowTargetHandleId = uuidv4();
+const subworkflowOutputNodeId = uuidv4();
+const subworkflowOutputId = uuidv4();
+
+export default {
+  input_variables: [],
+  output_variables: [
+    {
+      id: outputVariableId,
+      key: "result",
+      type: "STRING",
+    },
+  ],
+  state_variables: [
+    {
+      id: parentStateVariableId,
+      key: "parent_counter",
+      type: "NUMBER",
+      default: {
+        type: "NUMBER",
+        value: 0,
+      },
+    },
+  ],
+  workflow_raw_data: {
+    definition: {
+      module: ["test", "workflow"],
+      name: "Workflow",
+    },
+    edges: [
+      {
+        id: uuidv4(),
+        source_handle_id: entrypointHandleId,
+        source_node_id: entrypointNodeId,
+        target_handle_id: subworkflowTargetHandleId,
+        target_node_id: subworkflowNodeId,
+        type: "DEFAULT",
+      },
+      {
+        id: uuidv4(),
+        source_handle_id: subworkflowSourceHandleId,
+        source_node_id: subworkflowNodeId,
+        target_handle_id: outputHandleId,
+        target_node_id: outputNodeId,
+        type: "DEFAULT",
+      },
+    ],
+    nodes: [
+      {
+        base: null,
+        data: {
+          label: "Entrypoint Node",
+          source_handle_id: entrypointHandleId,
+        },
+        definition: null,
+        id: entrypointNodeId,
+        inputs: [],
+        type: "ENTRYPOINT",
+      },
+      {
+        base: {
+          module: [
+            "vellum",
+            "workflows",
+            "nodes",
+            "displayable",
+            "inline_subworkflow_node",
+            "node",
+          ],
+          name: "InlineSubworkflowNode",
+        },
+        data: {
+          label: "Subworkflow With State",
+          error_output_id: null,
+          input_variables: [],
+          output_variables: [
+            {
+              id: subworkflowOutputId,
+              key: "subworkflow_result",
+              type: "STRING",
+            },
+          ],
+          state_variables: [
+            {
+              id: innerStateVariableId,
+              key: "inner_counter",
+              type: "NUMBER",
+              default: {
+                type: "NUMBER",
+                value: 0,
+              },
+            },
+          ],
+          source_handle_id: subworkflowSourceHandleId,
+          target_handle_id: subworkflowTargetHandleId,
+          variant: "INLINE",
+          workflow_raw_data: {
+            definition: {
+              module: ["test", "workflow", "subworkflow_with_state"],
+              name: "SubworkflowWithState",
+            },
+            edges: [
+              {
+                id: uuidv4(),
+                source_handle_id: "inner-entry-src",
+                source_node_id: "inner-entry-id",
+                target_handle_id: "inner-final-handle",
+                target_node_id: subworkflowOutputNodeId,
+                type: "DEFAULT",
+              },
+            ],
+            nodes: [
+              {
+                base: null,
+                data: {
+                  label: "Entrypoint Node",
+                  source_handle_id: "inner-entry-src",
+                },
+                definition: null,
+                id: "inner-entry-id",
+                inputs: [],
+                type: "ENTRYPOINT",
+              },
+              {
+                base: {
+                  module: [
+                    "vellum",
+                    "workflows",
+                    "nodes",
+                    "displayable",
+                    "final_output_node",
+                    "node",
+                  ],
+                  name: "FinalOutputNode",
+                },
+                data: {
+                  label: "Final Output",
+                  name: "subworkflow_result",
+                  node_input_id: "inner-output-input-id",
+                  output_id: subworkflowOutputId,
+                  output_type: "STRING",
+                  target_handle_id: "inner-final-handle",
+                },
+                definition: {
+                  module: [
+                    "test",
+                    "workflow",
+                    "subworkflow_with_state",
+                    "final_output",
+                  ],
+                  name: "SubworkflowFinalOutput",
+                },
+                id: subworkflowOutputNodeId,
+                inputs: [
+                  {
+                    id: "inner-output-input-id",
+                    key: "node_input",
+                    value: {
+                      combinator: "OR",
+                      rules: [
+                        {
+                          data: {
+                            type: "STRING",
+                            value: "result from subworkflow",
+                          },
+                          type: "CONSTANT_VALUE",
+                        },
+                      ],
+                    },
+                  },
+                ],
+                outputs: [
+                  {
+                    id: subworkflowOutputId,
+                    name: "value",
+                    type: "STRING",
+                    value: {
+                      type: "CONSTANT_VALUE",
+                      value: {
+                        type: "STRING",
+                        value: "result from subworkflow",
+                      },
+                    },
+                  },
+                ],
+                ports: [],
+                trigger: {
+                  id: "inner-final-handle",
+                  merge_behavior: "AWAIT_ANY",
+                },
+                type: "TERMINAL",
+              },
+            ],
+            output_values: [
+              {
+                output_variable_id: subworkflowOutputId,
+                value: {
+                  node_id: subworkflowOutputNodeId,
+                  node_output_id: subworkflowOutputId,
+                  type: "NODE_OUTPUT",
+                },
+              },
+            ],
+          },
+        },
+        definition: {
+          module: ["test", "workflow", "subworkflow_with_state"],
+          name: "SubworkflowWithStateNode",
+        },
+        id: subworkflowNodeId,
+        inputs: [],
+        ports: [
+          {
+            id: subworkflowSourceHandleId,
+            name: "default",
+            type: "DEFAULT",
+          },
+        ],
+        trigger: {
+          id: subworkflowTargetHandleId,
+          merge_behavior: "AWAIT_ATTRIBUTES",
+        },
+        type: "SUBWORKFLOW",
+      },
+      {
+        base: {
+          module: [
+            "vellum",
+            "workflows",
+            "nodes",
+            "displayable",
+            "final_output_node",
+            "node",
+          ],
+          name: "FinalOutputNode",
+        },
+        data: {
+          label: "Workflow Output",
+          name: "result",
+          node_input_id: outputInputId,
+          output_id: outputVariableId,
+          output_type: "STRING",
+          target_handle_id: outputHandleId,
+        },
+        definition: {
+          module: ["test", "workflow", "final_output"],
+          name: "WorkflowFinalOutput",
+        },
+        id: outputNodeId,
+        inputs: [
+          {
+            id: outputInputId,
+            key: "node_input",
+            value: {
+              combinator: "OR",
+              rules: [
+                {
+                  data: {
+                    node_id: subworkflowNodeId,
+                    output_id: subworkflowOutputId,
+                  },
+                  type: "NODE_OUTPUT",
+                },
+              ],
+            },
+          },
+        ],
+        outputs: [
+          {
+            id: outputVariableId,
+            name: "value",
+            type: "STRING",
+            value: {
+              node_id: subworkflowNodeId,
+              node_output_id: subworkflowOutputId,
+              type: "NODE_OUTPUT",
+            },
+          },
+        ],
+        ports: [],
+        trigger: {
+          id: outputHandleId,
+          merge_behavior: "AWAIT_ANY",
+        },
+        type: "TERMINAL",
+      },
+    ],
+    output_values: [
+      {
+        output_variable_id: outputVariableId,
+        value: {
+          node_id: outputNodeId,
+          node_output_id: outputVariableId,
+          type: "NODE_OUTPUT",
+        },
+      },
+    ],
+  },
+  assertions: ["state.py", "nodes/subworkflow_with_state/state.py"],
+};

ee/codegen/src/generators/nodes/inline-subworkflow-node.ts

@@ -128,7 +128,7 @@ export class InlineSubworkflowNode extends BaseNestedWorkflowNode<
       workflowVersionExecConfig: {
         workflowRawData: inlineSubworkflowNodeData.workflowRawData,
         inputVariables: inlineSubworkflowNodeData.inputVariables,
-        stateVariables: [],
+        stateVariables: inlineSubworkflowNodeData.stateVariables ?? [],
         outputVariables: inlineSubworkflowNodeData.outputVariables,
       },
       moduleName: nestedWorkflowContext.moduleName,

ee/codegen/src/serializers/vellum.ts

@@ -1239,6 +1239,10 @@ export const InlineSubworkflowNodeDataSerializer: ObjectSchema<
     "output_variables",
     listSchema(VellumVariableSerializer)
   ),
+  stateVariables: propertySchema(
+    "state_variables",
+    listSchema(VellumVariableSerializer).optional()
+  ),
   label: stringSchema(),
   sourceHandleId: propertySchema("source_handle_id", stringSchema()),
   targetHandleId: propertySchema("target_handle_id", stringSchema()),
@@ -1250,6 +1254,7 @@ export declare namespace InlineSubworkflowNodeDataSerializer {
     workflow_raw_data: WorkflowRawDataSerializer.Raw;
     input_variables: VellumVariableSerializer.Raw[];
     output_variables: VellumVariableSerializer.Raw[];
+    state_variables?: VellumVariableSerializer.Raw[] | null;
     label: string;
     source_handle_id: string;
     target_handle_id: string;

ee/codegen/src/types/vellum.ts

@@ -522,6 +522,7 @@ export interface InlineSubworkflowNodeData {
   workflowRawData: WorkflowRawData;
   inputVariables: VellumVariable[];
   outputVariables: VellumVariable[];
+  stateVariables?: VellumVariable[];
   label: string;
   sourceHandleId: string;
   targetHandleId: string;