Add configurable state reference to ChatMessageTrigger

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

@@ -72,6 +72,45 @@ class Workflow(BaseWorkflow):
 "
 `;
 
+exports[`generateCode > should generate code for %1 chat-message-trigger-custom-key.ts > triggers/custom_state_trigger.py 1`] = `
+"from vellum.workflows.references import LazyReference
+from vellum.workflows.triggers import ChatMessageTrigger
+
+from ..state import State
+
+
+class CustomStateTrigger(ChatMessageTrigger):
+    class Config(ChatMessageTrigger.Config):
+        output = LazyReference("Workflow.Outputs.result")
+        state = State.messages
+
+    class Display(ChatMessageTrigger.Display):
+        label = "Custom State Trigger"
+        x = 100
+        y = 200
+        z_index = 1
+        icon = "vellum:icon:message"
+        color = "blue"
+"
+`;
+
+exports[`generateCode > should generate code for %1 chat-message-trigger-custom-key.ts > workflow.py 1`] = `
+"from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs import BaseInputs
+
+from .nodes.some_node import SomeNode
+from .state import State
+from .triggers.custom_state_trigger import CustomStateTrigger
+
+
+class Workflow(BaseWorkflow[BaseInputs, State]):
+    graph = CustomStateTrigger >> SomeNode
+
+    class Outputs(BaseWorkflow.Outputs):
+        result = SomeNode.Outputs.result
+"
+`;
+
 exports[`generateCode > should generate code for %1 code-execution-node-with-await-all.ts > nodes/code_execution_with_await_all/__init__.py 1`] = `
 "from vellum.workflows.nodes.displayable import CodeExecutionNode
 from vellum.workflows.state import BaseState

ee/codegen/src/__test__/generate-code-fixtures/chat-message-trigger-custom-key.ts

@@ -0,0 +1,107 @@
+export default {
+  workflow_raw_data: {
+    nodes: [
+      {
+        id: "some-node",
+        type: "GENERIC",
+        label: "Some Node",
+        display_data: null,
+        base: {
+          name: "BaseNode",
+          module: ["vellum", "workflows", "nodes", "bases", "base"],
+        },
+        definition: {
+          name: "SomeNode",
+          module: ["testing", "nodes", "some_node"],
+        },
+        trigger: {
+          id: "some-target",
+          merge_behavior: "AWAIT_ATTRIBUTES",
+        },
+        ports: [
+          {
+            id: "some-default-port-id",
+            name: "default",
+            type: "DEFAULT",
+          },
+        ],
+        outputs: [
+          {
+            id: "some-output-id",
+            name: "result",
+            type: "STRING",
+          },
+        ],
+        attributes: [],
+      },
+    ],
+    edges: [
+      {
+        id: "edge-2",
+        source_node_id: "chat-message-trigger",
+        source_handle_id: "chat-message-trigger",
+        target_node_id: "some-node",
+        target_handle_id: "some-target",
+        type: "DEFAULT",
+      },
+    ],
+    output_values: [
+      {
+        output_variable_id: "workflow-output-variable-id",
+        value: {
+          type: "NODE_OUTPUT",
+          node_id: "some-node",
+          node_output_id: "some-output-id",
+        },
+      },
+    ],
+  },
+  input_variables: [],
+  output_variables: [
+    {
+      id: "workflow-output-variable-id",
+      key: "result",
+      type: "STRING",
+    },
+  ],
+  state_variables: [
+    {
+      id: "messages-state-variable-id",
+      key: "messages",
+      type: "CHAT_HISTORY",
+    },
+  ],
+  triggers: [
+    {
+      id: "chat-message-trigger",
+      type: "CHAT_MESSAGE",
+      attributes: [
+        {
+          id: "message-attribute-id",
+          key: "message",
+          type: "JSON",
+        },
+      ],
+      exec_config: {
+        output: {
+          type: "WORKFLOW_OUTPUT",
+          output_variable_id: "workflow-output-variable-id",
+        },
+        state: {
+          state_variable_id: "messages-state-variable-id",
+        },
+      },
+      display_data: {
+        label: "Custom State Trigger",
+        position: {
+          x: 100,
+          y: 200,
+        },
+        z_index: 1,
+        icon: "vellum:icon:message",
+        color: "blue",
+      },
+    },
+  ],
+  assertions: ["workflow.py", "triggers/custom_state_trigger.py"],
+};

ee/codegen/src/generators/triggers/chat-message-trigger.ts

@@ -12,6 +12,7 @@ import type { AstNode } from "src/generators/extensions/ast-node";
 import type {
   ChatMessageTrigger as ChatMessageTriggerType,
   WorkflowOutputWorkflowReference,
+  WorkflowStateVariableWorkflowReference,
 } from "src/types/vellum";
 
 export declare namespace ChatMessageTriggerGenerator {
@@ -46,17 +47,24 @@ export class ChatMessageTrigger extends BaseTrigger<ChatMessageTriggerType> {
   protected getTriggerClassBody(): AstNode[] {
     const body: AstNode[] = [];
 
-    body.push(...this.createAttributeFields());
+    // Note: We don't call createAttributeFields() here because the `message` attribute
+    // is inherited from the Python ChatMessageTrigger base class
 
     const execConfig = this.trigger.execConfig;
-    if (execConfig?.output) {
-      body.push(this.createConfigClass(execConfig.output));
+    const hasOutput = execConfig?.output;
+    const hasState = execConfig?.state;
+
+    if (hasOutput || hasState) {
+      body.push(this.createConfigClass(execConfig?.output, execConfig?.state));
     }
 
     return body;
   }
 
-  private createConfigClass(output: WorkflowOutputWorkflowReference): AstNode {
+  private createConfigClass(
+    output?: WorkflowOutputWorkflowReference,
+    state?: Omit<WorkflowStateVariableWorkflowReference, "type">
+  ): AstNode {
     const configClass = new Class({
       name: "Config",
       extends_: [
@@ -68,14 +76,47 @@ export class ChatMessageTrigger extends BaseTrigger<ChatMessageTriggerType> {
       ],
     });
 
-    const outputField = this.createOutputField(output);
-    if (outputField) {
-      configClass.add(outputField);
+    if (output) {
+      const outputField = this.createOutputField(output);
+      if (outputField) {
+        configClass.add(outputField);
+      }
+    }
+
+    if (state) {
+      const stateField = this.createStateField(state);
+      if (stateField) {
+        configClass.add(stateField);
+      }
     }
 
     return configClass;
   }
 
+  private createStateField(
+    state: Omit<WorkflowStateVariableWorkflowReference, "type">
+  ): AstNode | undefined {
+    const stateVariableContext =
+      this.workflowContext.findStateVariableContextById(state.stateVariableId);
+
+    if (!stateVariableContext) {
+      return undefined;
+    }
+
+    const stateReference = new Reference({
+      name: "State",
+      modulePath: stateVariableContext.definition.module,
+      attribute: [stateVariableContext.name],
+    });
+
+    this.inheritReferences(stateReference);
+
+    return new Field({
+      name: "state",
+      initializer: stateReference,
+    });
+  }
+
   private createOutputField(
     output: WorkflowOutputWorkflowReference
   ): AstNode | undefined {

ee/codegen/src/serializers/vellum.ts

@@ -2453,11 +2453,13 @@ export declare namespace IntegrationTriggerSerializer {
 
 const ChatMessageTriggerExecConfigSerializer = objectSchema({
   output: WorkflowValueDescriptorSerializer.optional(),
+  state: WorkflowStateVariableWorkflowReferenceSerializer.optional(),
 });
 
 export declare namespace ChatMessageTriggerExecConfigSerializer {
   interface Raw {
     output?: WorkflowValueDescriptorSerializer.Raw | null;
+    state?: WorkflowStateVariableWorkflowReferenceSerializer.Raw | null;
   }
 }
 

ee/codegen/src/types/vellum.ts

@@ -852,6 +852,7 @@ export interface IntegrationTrigger extends BaseTrigger {
 
 export interface ChatMessageTriggerExecConfig {
   output?: WorkflowOutputWorkflowReference;
+  state?: Omit<WorkflowStateVariableWorkflowReference, "type">;
 }
 
 export interface ChatMessageTrigger extends BaseTrigger {

ee/vellum_ee/workflows/display/tests/workflow_serialization/test_chat_message_trigger_custom_state_serialization.py

@@ -0,0 +1,69 @@
+"""Tests for ChatMessageTrigger with custom state reference serialization."""
+
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.chat_message_trigger_custom_state.workflows.custom_state_workflow import CustomStateWorkflow
+
+
+def test_custom_state_workflow_serialization():
+    """CustomStateWorkflow serializes correctly with ChatMessageTrigger using custom state reference."""
+
+    # GIVEN a Workflow that uses a ChatMessageTrigger with custom state reference
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=CustomStateWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+        "triggers",
+    }
+
+    # AND the triggers should be serialized correctly with custom state reference
+    assert serialized_workflow["triggers"] == [
+        {
+            "id": "ee5cb788-8e76-4ddb-af9b-37d3b176acde",
+            "type": "CHAT_MESSAGE",
+            "attributes": [
+                {
+                    "id": "6b15244a-d2ee-450f-a702-678e67c62b4a",
+                    "key": "message",
+                    "type": "ARRAY",
+                    "required": True,
+                    "default": {
+                        "type": "ARRAY",
+                        "value": None,
+                    },
+                    "extensions": None,
+                    "schema": None,
+                }
+            ],
+            "exec_config": {
+                "output": {
+                    "type": "WORKFLOW_OUTPUT",
+                    "output_variable_id": "02ea4c1b-26a0-4dbb-b9f0-2fad755b794f",
+                },
+                "state": {
+                    "state_variable_id": "8b859454-fd27-4e7b-9436-cfb4fca6fdcc",
+                },
+            },
+            "display_data": {
+                "label": "Chat Message",
+                "position": {
+                    "x": 0.0,
+                    "y": 0.0,
+                },
+                "z_index": 0,
+                "icon": "vellum:icon:message-dots",
+                "color": "blue",
+            },
+        }
+    ]
+
+    # AND the state_variable_id should match the corresponding state variable
+    state_variable_id = serialized_workflow["triggers"][0]["exec_config"]["state"]["state_variable_id"]
+    state_variable_ids = [sv["id"] for sv in serialized_workflow["state_variables"]]
+    assert state_variable_id in state_variable_ids