vellum-workflow-server 1.4.2__py3-none-any.whl → 1.4.2.post2__py3-none-any.whl

{vellum_workflow_server-1.4.2.dist-info → vellum_workflow_server-1.4.2.post2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 1.4.2
+Version: 1.4.2.post2
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4

{vellum_workflow_server-1.4.2.dist-info → vellum_workflow_server-1.4.2.post2.dist-info}/RECORD

@@ -4,17 +4,17 @@ workflow_server/api/auth_middleware.py,sha256=IlZaCiwZ5nwQqk5sYQorvOFj7lt0p1ZSSE
 workflow_server/api/healthz_view.py,sha256=itiRvBDBXncrw8Kbbc73UZLwqMAhgHOR3uSre_dAfgY,404
 workflow_server/api/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/api/tests/test_input_display_mapping.py,sha256=drBZqMudFyB5wgiUOcMgRXz7E7ge-Qgxbstw4E4f0zE,2211
-workflow_server/api/tests/test_workflow_view.py,sha256=RlAw1tHeIlnOXGrFQN-w3EOLPZkhp6Dfy6d1r7kU5oc,22573
+workflow_server/api/tests/test_workflow_view.py,sha256=d46UNMZUJbIKLiTJkjIsVpgqCJMwCe4LL6RjFkfx_Y4,29178
 workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=FaEIgGsbq8M7ZF2opVjBdGCYeRPF_vqsUhFTruLInxA,31358
-workflow_server/api/workflow_view.py,sha256=UZTxWImM9kmof819SCa3ljJKiYEsCrBFD8vp2_f-zAg,21719
+workflow_server/api/workflow_view.py,sha256=qayjwq18DSbmuMa6ROksc68v4XLcQOr7hAH3yRV5mlQ,21443
 workflow_server/code_exec_runner.py,sha256=E-HsjAL53L-znSMPg7lDiQNzyCjL6W076ZoWWbrSRrU,2217
 workflow_server/config.py,sha256=qmmTr6ty3ZN5LDOFs3TfUxYshYe6Mmn_LanplHHeE9Q,1796
 workflow_server/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/core/cancel_workflow.py,sha256=Ffkc3mzmrdMEUcD-sHfEhX4IwVrka-E--SxKA1dUfIU,2185
 workflow_server/core/events.py,sha256=24MA66DVQuaLJJcZrS8IL1Zq4Ohi9CoouKZ5VgoH3Cs,1402
-workflow_server/core/executor.py,sha256=5cji5KQSukLrABmihJp9cgKjXS145TocOz2cOcnof04,18962
-workflow_server/core/utils.py,sha256=cmwHbKCfXqtUutBD3akGus0Ga7a1xG3zlOw-jEMx6mI,1795
-workflow_server/core/workflow_executor_context.py,sha256=VafZg74t_GQ3_2DEWVroy38rSy_spcAw4c3NrOOWOKY,3198
+workflow_server/core/executor.py,sha256=zr5iruCsAMI5yHoW_yWzAhCo1U4Si9Kr9tjOmBKQf4A,20102
+workflow_server/core/utils.py,sha256=si0NB4Suurc-mn8NYdn59xM9CkPrfOP1aWEVrZvifDI,1929
+workflow_server/core/workflow_executor_context.py,sha256=w0nFhtu54GX36ZzmaqLTD3-ssPtoFd9QX02dPgeIbBM,3615
 workflow_server/server.py,sha256=QBU12AaAfAgLqfCDBd24qIJl_mbheiq0-hfcWV7rZM4,1234
 workflow_server/start.py,sha256=pkwRcms6I4tkVHP06LdrZY6rG_DFHfBx4ioY5X91W5k,2264
 workflow_server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -28,7 +28,7 @@ workflow_server/utils/tests/test_sentry_integration.py,sha256=LGmWiaLhFrx-jslrRj
 workflow_server/utils/tests/test_system_utils.py,sha256=_4GwXvVvU5BrATxUEWwQIPg0bzQXMWBtiBmjP8MTxJM,4314
 workflow_server/utils/tests/test_utils.py,sha256=0Nq6du8o-iBtTrip9_wgHES53JSiJbVdSXaBnPobw3s,6930
 workflow_server/utils/utils.py,sha256=ZPoM1Suhid22dpB8oEFLux8wx-9iyzmSfWuYxSCrgWk,4774
-vellum_workflow_server-1.4.2.dist-info/METADATA,sha256=ZevXg46YiyUvi73O9rmNa3tYSRfKL1JPRM5wKH-ypxs,2267
-vellum_workflow_server-1.4.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_workflow_server-1.4.2.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
-vellum_workflow_server-1.4.2.dist-info/RECORD,,
+vellum_workflow_server-1.4.2.post2.dist-info/METADATA,sha256=mfiAFH3GV0WI9qDRFRCIwIHE3e8DcxZfbhvOIlJWpSM,2273
+vellum_workflow_server-1.4.2.post2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-1.4.2.post2.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-1.4.2.post2.dist-info/RECORD,,

workflow_server/api/tests/test_workflow_view.py

@@ -1,7 +1,8 @@
+import json
 import logging
 import re
 from unittest.mock import patch
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from deepdiff import DeepDiff
 
@@ -534,3 +535,197 @@ def test_serialize_route__with_invalid_workspace_api_key():
 
     # AND the response should contain the serialization result
     assert "exec_config" in response.json
+
+
+def test_stream_node_route__with_node_id():
+    """
+    Tests that the stream-node endpoint works with node_id.
+    """
+    # GIVEN a valid request body with node_id
+    node_id = uuid4()
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "node_id": str(node_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "from .display import *",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.test_node import TestNode
+
+class Workflow(BaseWorkflow):
+    graph = TestNode
+
+    class Outputs(BaseWorkflow.Outputs):
+        result = TestNode.Outputs.value
+""",
+            "nodes/__init__.py": "from .test_node import TestNode\n__all__ = ['TestNode']",
+            "nodes/test_node.py": """\
+from vellum.workflows.nodes import BaseNode
+
+class TestNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        value = "test_result"
+""",
+            "display/__init__.py": "from .nodes import *\nfrom .workflow import *",
+            "display/workflow.py": """\
+from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+""",
+            "display/nodes/__init__.py": "from .test_node import TestNodeDisplay\n__all__ = ['TestNodeDisplay']",
+            "display/nodes/test_node.py": f"""\
+from uuid import UUID
+from vellum_ee.workflows.display.nodes import BaseNodeDisplay
+from ...nodes.test_node import TestNode
+
+class TestNodeDisplay(BaseNodeDisplay[TestNode]):
+    node_id = UUID("{node_id}")
+""",
+        },
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 200 response
+    assert response.status_code == 200
+
+    events = [json.loads(line) for line in response.data.decode().strip().split("\n") if line.strip()]
+    assert len(events) == 4
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "node.execution.initiated"
+    assert events[2]["name"] == "node.execution.fulfilled"
+    assert events[3]["name"] == "vembda.execution.fulfilled"
+
+
+def test_stream_node_route__with_node_module_and_name_backward_compatibility():
+    """
+    Tests that the stream-node endpoint still works with node_module and node_name for backward compatibility.
+    """
+    # GIVEN a valid request body with node_module and node_name (old format)
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "node_module": "nodes.test_node",
+        "node_name": "TestNode",
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.test_node import TestNode
+
+class Workflow(BaseWorkflow):
+    graph = TestNode
+""",
+            "nodes/__init__.py": "from .test_node import TestNode",
+            "nodes/test_node.py": """\
+from vellum.workflows.nodes import BaseNode
+
+class TestNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        value = "test_result"
+""",
+        },
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 200 response
+    assert response.status_code == 200
+
+    events = [json.loads(line) for line in response.data.decode().strip().split("\n") if line.strip()]
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "node.execution.initiated", json.dumps(events[1]["body"])
+    assert events[2]["name"] == "node.execution.fulfilled"
+    assert events[3]["name"] == "vembda.execution.fulfilled"
+    assert len(events) == 4
+
+
+def test_stream_node_route__missing_node_info_validation():
+    """
+    Tests that the stream-node endpoint returns validation error when neither
+    node_id nor node_module/node_name are provided.
+    """
+    # GIVEN a request body missing node identification
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {"__init__.py": "", "workflow.py": ""},
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 400 response
+    assert response.status_code == 400
+
+    # AND we get a validation error message
+    assert "Either node_id or both node_module and node_name must be provided" in response.get_json()["detail"]
+
+
+def test_stream_node_route__invalid_node_id():
+    """
+    Tests that the stream-node endpoint returns 404 for invalid node_id.
+    """
+    # GIVEN a request body with invalid node_id
+    invalid_node_id = uuid4()
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "node_id": str(invalid_node_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow):
+    pass
+""",
+        },
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 200 response
+    # TODO: In the future, we would want this to return a 4xx response by returning the workflow
+    # instance and the node definition that we want to run as part of request deserialization.
+    assert response.status_code == 200, response.text
+
+    # AND we get an appropriate error message
+    # TODO: In a future where we are returning 4xx responses, we assert the following data:
+    # response_data = response.get_json()
+    # assert "Node with ID" in response_data["detail"]
+    # assert "not found" in response_data["detail"]
+    events = [json.loads(line) for line in response.data.decode().strip().split("\n") if line.strip()]
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "vembda.execution.fulfilled"
+    assert len(events) == 2

workflow_server/api/workflow_view.py

@@ -36,7 +36,7 @@ from workflow_server.core.events import (
     VembdaExecutionInitiatedBody,
     VembdaExecutionInitiatedEvent,
 )
-from workflow_server.core.executor import stream_node_pebble_timeout, stream_workflow, stream_workflow_process_timeout
+from workflow_server.core.executor import stream_node_process_timeout, stream_workflow, stream_workflow_process_timeout
 from workflow_server.core.utils import (
     create_vembda_rejected_event,
     is_events_emitting_enabled,
@@ -374,10 +374,10 @@ def stream_node_route() -> Response:
         parent=None,
     )
 
-    app.logger.debug(f"Node stream received {data.get('execution_id')}")
+    app.logger.debug(f"Node stream started. Trace ID: {context.trace_id}")
 
     pebble_queue: Queue[dict] = Queue()
-    stream_future = stream_node_pebble_timeout(
+    process = stream_node_process_timeout(
         executor_context=context,
         queue=pebble_queue,
     )
@@ -388,14 +388,9 @@ def stream_node_route() -> Response:
                 event = pebble_queue.get(timeout=context.timeout)
 
             except Empty:
-                if stream_future.exception() is not None:
-                    # This happens when theres a problem with the stream function call
-                    # itself not the workflow runner
+                if not process.is_alive():
                     yield create_vembda_rejected_event(context, "Internal Server Error")
-                    app.logger.exception(stream_future.exception())
                     break
-                else:
-                    continue
 
             yield event
             if event.get("name") == VEMBDA_EXECUTION_FULFILLED_EVENT_NAME:

workflow_server/core/executor.py

@@ -15,7 +15,6 @@ from traceback import format_exc
 from uuid import UUID, uuid4
 from typing import Any, Callable, Generator, Iterator, Optional, Tuple, Type
 
-from pebble import concurrent
 from vellum_ee.workflows.display.utils.events import event_enricher
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
@@ -24,6 +23,14 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.emitters.vellum_emitter import VellumEmitter
 from vellum.workflows.events.exception_handling import stream_initialization_exception
+from vellum.workflows.events.node import (
+    NodeExecutionFulfilledBody,
+    NodeExecutionFulfilledEvent,
+    NodeExecutionInitiatedBody,
+    NodeExecutionInitiatedEvent,
+    NodeExecutionStreamingBody,
+    NodeExecutionStreamingEvent,
+)
 from vellum.workflows.events.types import BaseEvent
 from vellum.workflows.events.workflow import WorkflowEventDisplayContext
 from vellum.workflows.exceptions import WorkflowInitializationException
@@ -44,9 +51,12 @@ from workflow_server.core.events import (
     VembdaExecutionFulfilledBody,
     VembdaExecutionFulfilledEvent,
 )
-from workflow_server.core.utils import is_events_emitting_enabled, serialize_vembda_rejected_event
+from workflow_server.core.utils import (
+    create_vembda_rejected_event,
+    is_events_emitting_enabled,
+    serialize_vembda_rejected_event,
+)
 from workflow_server.core.workflow_executor_context import (
-    DEFAULT_TIMEOUT_SECONDS,
     BaseExecutorContext,
     NodeExecutorContext,
     WorkflowExecutorContext,
@@ -56,38 +66,31 @@ from workflow_server.utils.log_proxy import redirect_log
 logger = logging.getLogger(__name__)
 
 
-@concurrent.process(timeout=DEFAULT_TIMEOUT_SECONDS)
-# type ignore since pebble annotation changes return type
-def stream_node_pebble_timeout(
+def stream_node_process_timeout(
     executor_context: NodeExecutorContext,
     queue: Queue,
-) -> None:
-    _stream_node_wrapper(
-        executor_context=executor_context,
-        queue=queue,
+) -> Process:
+    node_process = Process(
+        target=_stream_node_wrapper,
+        args=(executor_context, queue),
     )
+    node_process.start()
+
+    if node_process.exitcode is not None:
+        queue.put(create_vembda_rejected_event(executor_context, "Internal Server Error", timed_out=True))
+
+    return node_process
 
 
 def _stream_node_wrapper(executor_context: NodeExecutorContext, queue: Queue) -> None:
     try:
         for event in stream_node(executor_context=executor_context):
             queue.put(event)
+    except WorkflowInitializationException as e:
+        queue.put(create_vembda_rejected_event(executor_context, e.message))
     except Exception as e:
         logger.exception(e)
-        queue.put(
-            VembdaExecutionFulfilledEvent(
-                id=uuid4(),
-                timestamp=datetime.now(),
-                trace_id=executor_context.trace_id,
-                span_id=executor_context.execution_id,
-                body=VembdaExecutionFulfilledBody(
-                    exit_code=-1,
-                    stderr="Internal Server Error",
-                    container_overhead_latency=executor_context.container_overhead_latency,
-                ),
-                parent=None,
-            ).model_dump(mode="json")
-        )
+        queue.put(create_vembda_rejected_event(executor_context, "Internal Server Error"))
 
 
 def _stream_workflow_wrapper(
@@ -125,7 +128,7 @@ def _stream_workflow_wrapper(
 def stream_workflow_process_timeout(
     executor_context: WorkflowExecutorContext,
     queue: Queue,
-    cancel_signal: Optional[ThreadingEvent],
+    cancel_signal: ThreadingEvent,
     timeout_signal: ThreadingEvent,
 ) -> Process:
     workflow_process = Process(
@@ -140,19 +143,7 @@ def stream_workflow_process_timeout(
     workflow_process.start()
 
     if workflow_process.exitcode is not None:
-        vembda_fulfilled_event = VembdaExecutionFulfilledEvent(
-            id=uuid4(),
-            timestamp=datetime.now(),
-            trace_id=executor_context.trace_id,
-            span_id=executor_context.execution_id,
-            body=VembdaExecutionFulfilledBody(
-                exit_code=-1,
-                timed_out=True,
-                container_overhead_latency=executor_context.container_overhead_latency,
-            ),
-            parent=None,
-        )
-        queue.put(vembda_fulfilled_event.model_dump(mode="json"))
+        queue.put(create_vembda_rejected_event(executor_context, "Internal Server Error", timed_out=True))
 
     return workflow_process
 
@@ -259,9 +250,7 @@ def stream_workflow(
         if emitter_thread:
             emitter_thread.join()
 
-        timer_thread = next((t for t in threading.enumerate() if t.name.startswith("Thread-")), None)
-        if timer_thread:
-            timer_thread.join()
+        workflow.join()
 
     return (
         _call_stream(
@@ -278,28 +267,74 @@ def stream_node(
     executor_context: NodeExecutorContext,
     disable_redirect: bool = True,
 ) -> Iterator[dict]:
-    def call_node() -> Generator[dict[str, Any], Any, None]:
-        workflow, namespace = _gather_workflow(executor_context)
-        node_module = importlib.import_module(f"{namespace}.{executor_context.node_module}")
-
-        Node = getattr(node_module, executor_context.node_name)
-
-        workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
-        workflow_state = _get_workflow_state(executor_context, workflow_inputs=workflow_inputs)
-
-        node = Node(
-            state=workflow_state,
-            context=workflow._context,
+    workflow, namespace = _gather_workflow(executor_context)
+    Node: Optional[Type[BaseNode]] = None
+
+    for workflow_node in workflow.get_nodes():
+        if executor_context.node_id and workflow_node.__id__ == executor_context.node_id:
+            Node = workflow_node
+            break
+        elif (
+            executor_context.node_module
+            and executor_context.node_name
+            and workflow_node.__name__ == executor_context.node_name
+            and workflow_node.__module__ == f"{namespace}.{executor_context.node_module}"
+        ):
+            Node = workflow_node
+            break
+
+    if not Node:
+        identifier = executor_context.node_id or f"{executor_context.node_module}.{executor_context.node_name}"
+        raise WorkflowInitializationException(
+            message=f"Node '{identifier}' not found in workflow",
+            workflow_definition=workflow.__class__,
         )
 
+    workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
+    workflow_state = _get_workflow_state(executor_context, workflow_inputs=workflow_inputs)
+    node = Node(
+        state=workflow_state,
+        context=workflow._context,
+    )
+
+    def call_node() -> Generator[dict[str, Any], Any, None]:
         executor_context.stream_start_time = time.time_ns()
+        span_id = uuid4()
+        yield NodeExecutionInitiatedEvent(
+            id=uuid4(),
+            timestamp=datetime.now(),
+            trace_id=executor_context.trace_id,
+            span_id=span_id,
+            body=NodeExecutionInitiatedBody(
+                node_definition=Node,
+                inputs=node._inputs,
+            ),
+        ).model_dump(mode="json")
         node_outputs = node.run()
 
         if isinstance(node_outputs, (Iterator)):
             for node_output in node_outputs:
-                yield json.loads(json.dumps(node_output, default=vars))
+                yield NodeExecutionStreamingEvent(
+                    id=uuid4(),
+                    timestamp=datetime.now(),
+                    trace_id=executor_context.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionStreamingBody(
+                        node_definition=Node,
+                        output=node_output,
+                    ),
+                ).model_dump(mode="json")
         else:
-            yield json.loads(json.dumps(node_outputs, default=vars))
+            yield NodeExecutionFulfilledEvent(
+                id=uuid4(),
+                timestamp=datetime.now(),
+                trace_id=executor_context.trace_id,
+                span_id=span_id,
+                body=NodeExecutionFulfilledBody(
+                    node_definition=Node,
+                    outputs=node_outputs,
+                ),
+            ).model_dump(mode="json")
 
     return _call_stream(
         executor_context=executor_context,

workflow_server/core/utils.py

@@ -7,7 +7,7 @@ from workflow_server.core.workflow_executor_context import BaseExecutorContext
 
 
 def _create_vembda_rejected_event_base(
-    executor_context: Optional[BaseExecutorContext], error_message: str
+    executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool
 ) -> VembdaExecutionFulfilledEvent:
     if executor_context:
         trace_id = executor_context.trace_id
@@ -25,6 +25,7 @@ def _create_vembda_rejected_event_base(
         span_id=span_id,
         body=VembdaExecutionFulfilledBody(
             exit_code=-1,
+            timed_out=timed_out,
             stderr=error_message,
             container_overhead_latency=container_overhead_latency,
         ),
@@ -32,12 +33,16 @@ def _create_vembda_rejected_event_base(
     )
 
 
-def create_vembda_rejected_event(executor_context: Optional[BaseExecutorContext], error_message: str) -> dict:
-    return _create_vembda_rejected_event_base(executor_context, error_message).model_dump(mode="json")
+def create_vembda_rejected_event(
+    executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool = False
+) -> dict:
+    return _create_vembda_rejected_event_base(executor_context, error_message, timed_out).model_dump(mode="json")
 
 
-def serialize_vembda_rejected_event(executor_context: Optional[BaseExecutorContext], error_message: str) -> str:
-    return _create_vembda_rejected_event_base(executor_context, error_message).model_dump_json()
+def serialize_vembda_rejected_event(
+    executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool = False
+) -> str:
+    return _create_vembda_rejected_event_base(executor_context, error_message, timed_out).model_dump_json()
 
 
 def is_events_emitting_enabled(executor_context: Optional[BaseExecutorContext]) -> bool:

workflow_server/core/workflow_executor_context.py

@@ -1,9 +1,11 @@
 from dataclasses import field
+from functools import cached_property
 import os
 from uuid import UUID
 from typing import Any, Optional
+from typing_extensions import Self
 
-from _pytest.compat import cached_property
+from pydantic import model_validator
 
 from vellum import ApiVersionEnum, Vellum, VellumEnvironment
 from vellum.client.core import UniversalBaseModel
@@ -95,5 +97,12 @@ class WorkflowExecutorContext(BaseExecutorContext):
 
 
 class NodeExecutorContext(BaseExecutorContext):
-    node_module: str
-    node_name: str
+    node_id: Optional[UUID] = None
+    node_module: Optional[str] = None
+    node_name: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_node_identification(self) -> Self:
+        if not self.node_id and not (self.node_module and self.node_name):
+            raise ValueError("Either node_id or both node_module and node_name must be provided")
+        return self