vellum-ai 1.11.2__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -19,6 +19,7 @@ from typing import (
     get_origin,
 )
 
+from vellum.client import Vellum as VellumClient
 from vellum.client.types.code_resource_definition import CodeResourceDefinition
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.constants import undefined
@@ -30,12 +31,14 @@ from vellum.workflows.references.node import NodeReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.functions import compile_annotation
 from vellum.workflows.utils.names import pascal_to_title_case
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
+from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.exceptions import NodeValidationError, UnsupportedSerializationException
 from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.registry import register_node_display_class
 
@@ -60,10 +63,9 @@ class BaseNodeDisplayMeta(type):
         cls = cast(Type["BaseNodeDisplay"], super().__new__(mcs, name, bases, dct))
         # This cast shouldn't be necessary, but it's a workaround for a mypy bug
         node_class = cast(Type[BaseNode], cls.infer_node_class() if name != "BaseNodeDisplay" else BaseNode)
-
         if not dct.get("output_display"):
             cls.output_display = {
-                ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name)
+                ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name, _is_implicit=True)
                 for ref in node_class.Outputs
                 if ref.name in node_class.__output_ids__
             }
@@ -140,6 +142,11 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     __unserializable_attributes__: Set[NodeReference] = set()
     # END: Attributes for backwards compatible serialization
 
+    def build(self, client: VellumClient) -> None:
+        # Individual display classes can override this method to perform any async logic
+        # needed to perform serialization across all nodes.
+        pass
+
     def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -202,14 +209,18 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         for port in node.Ports:
             id = str(self.get_node_port_display(port).id)
             if port._condition_type:
+                expression = None
+                if port._condition:
+                    try:
+                        expression = serialize_value(node_id, display_context, port._condition)
+                    except UnsupportedSerializationException as e:
+                        display_context.add_error(e)
                 ports.append(
                     {
                         "id": id,
                         "name": port.name,
                         "type": port._condition_type.value,
-                        "expression": (
-                            serialize_value(node_id, display_context, port._condition) if port._condition else None
-                        ),
+                        "expression": expression,
                     }
                 )
             else:
@@ -275,6 +286,16 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 if output in self.output_display
                 else str(uuid4_from_hash(f"{self.node_id}|{output.name}"))
             )
+            try:
+                schema = compile_annotation(output.normalized_type, {})
+            except Exception as e:
+                display_context.add_error(
+                    NodeValidationError(
+                        message=f"Failed to compile output schema for output '{output.name}': {e}",
+                        node_class_name=self._node.__name__,
+                    )
+                )
+                schema = None
 
             outputs.append(
                 {
@@ -282,6 +303,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                     "name": output.name,
                     "type": output_type,
                     "value": value,
+                    "schema": schema,
                 }
             )
 
@@ -480,6 +502,17 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         if self._node.Display.color is not None and self._node.Display.color != base_node.Display.color:
             base_kwargs["color"] = self._node.Display.color
 
+        # Add position from x, y if they exist
+        if self._node.Display.x is not None or self._node.Display.y is not None:
+            base_kwargs["position"] = NodeDisplayPosition(
+                x=self._node.Display.x if self._node.Display.x is not None else 0.0,
+                y=self._node.Display.y if self._node.Display.y is not None else 0.0,
+            )
+
+        # Add z_index if it exists
+        if self._node.Display.z_index is not None and self._node.Display.z_index != base_node.Display.z_index:
+            base_kwargs["z_index"] = self._node.Display.z_index
+
         # Add docstring as comment if present
         if docstring:
             base_kwargs["comment"] = NodeDisplayComment(value=docstring, expanded=True)
@@ -492,16 +525,16 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             # Get fields that were explicitly set for z_index handling
             fields_set = explicit_value.model_fields_set
 
+            # Override position only if explicitly set (since it has a default value)
+            if "position" in fields_set:
+                base_kwargs["position"] = explicit_value.position
+
             # Override simple attributes (only if not None)
-            for attr in ("position", "width", "height", "icon", "color"):
+            for attr in ("width", "height", "icon", "color", "z_index"):
                 value = getattr(explicit_value, attr, None)
                 if value is not None:
                     base_kwargs[attr] = value
 
-            # Include z_index if explicitly set (even if None)
-            if "z_index" in fields_set:
-                base_kwargs["z_index"] = explicit_value.z_index
-
             # Special handling for comment: merge docstring with explicit comment's expanded state
             if explicit_value.comment:
                 base_kwargs["comment"] = (

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -246,3 +246,96 @@ def test_serialize_display_class_used_as_fallback():
     assert display_data["icon"] == "vellum:icon:star"  # Falls back to Display class
     assert display_data["color"] == "green"  # Falls back to Display class
     assert display_data["position"] == {"x": 50, "y": 75}  # From BaseNodeDisplay
+
+
+def test_serialize_basenode_display_with_xyz():
+    """Tests that BaseNode.Display x, y, z_index attributes serialize correctly."""
+
+    # GIVEN a node with x, y, z_index in Display class
+    class MyNode(BaseNode):
+        class Display:
+            x = 100.0
+            y = 200.0
+            z_index = 5
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the display_data should include position and z_index from Display class
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["position"] == {"x": 100.0, "y": 200.0}
+    assert display_data["z_index"] == 5
+
+
+def test_serialize_basenode_display_with_partial_xy():
+    """Tests that BaseNode.Display with only x or y still creates position."""
+
+    # GIVEN a node with only x in Display class
+    class MyNodeX(BaseNode):
+        class Display:
+            x = 150.0
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyNodeX)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the display_data should include position with x and default y
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["position"] == {"x": 150.0, "y": 0.0}
+
+
+def test_serialize_explicit_display_data_overrides_display_class_xyz():
+    """Tests that BaseNodeDisplay's explicit display_data takes precedence over Display class x, y, z_index."""
+
+    # GIVEN a node with x, y, z_index in Display class
+    class MyNode(BaseNode):
+        class Display:
+            x = 100.0
+            y = 200.0
+            z_index = 5
+
+    # AND a BaseNodeDisplay with explicit display_data
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(
+            position=NodeDisplayPosition(x=300, y=400),
+            z_index=10,
+        )
+
+    # WHEN we serialize the node
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    # THEN the display_data should use values from BaseNodeDisplay
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["position"] == {"x": 300, "y": 400}  # BaseNodeDisplay overrides
+    assert display_data["z_index"] == 10  # BaseNodeDisplay overrides
+
+
+def test_serialize_display_class_xyz_used_as_fallback():
+    """Tests that Display class x, y, z_index are used when BaseNodeDisplay doesn't specify them."""
+
+    # GIVEN a node with x, y, z_index in Display class
+    class MyNode(BaseNode):
+        class Display:
+            x = 50.0
+            y = 75.0
+            z_index = 3
+
+    # AND a BaseNodeDisplay with only icon specified
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(
+            icon="vellum:icon:star",
+        )
+
+    # WHEN we serialize the node
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    # THEN the display_data should use x, y, z from Display class as fallback
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["position"] == {"x": 50.0, "y": 75.0}  # Falls back to Display class
+    assert display_data["z_index"] == 3  # Falls back to Display class
+    assert display_data["icon"] == "vellum:icon:star"  # From BaseNodeDisplay

vellum_ee/workflows/display/nodes/types.py

@@ -6,6 +6,7 @@ from uuid import UUID
 class NodeOutputDisplay:
     id: UUID
     name: str
+    _is_implicit: bool = False
 
 
 @dataclass

vellum_ee/workflows/display/nodes/vellum/__init__.py

@@ -4,7 +4,6 @@ from .code_execution_node import BaseCodeExecutionNodeDisplay
 from .conditional_node import BaseConditionalNodeDisplay
 from .error_node import BaseErrorNodeDisplay
 from .final_output_node import BaseFinalOutputNodeDisplay
-from .function_node import FunctionNodeDisplay
 from .guardrail_node import BaseGuardrailNodeDisplay
 from .inline_prompt_node import BaseInlinePromptNodeDisplay
 from .inline_subworkflow_node import BaseInlineSubworkflowNodeDisplay
@@ -26,7 +25,6 @@ __all__ = [
     "BaseConditionalNodeDisplay",
     "BaseErrorNodeDisplay",
     "BaseFinalOutputNodeDisplay",
-    "FunctionNodeDisplay",
     "BaseGuardrailNodeDisplay",
     "BaseInlinePromptNodeDisplay",
     "BaseInlineSubworkflowNodeDisplay",

vellum_ee/workflows/display/nodes/vellum/base_adornment_node.py

@@ -89,8 +89,11 @@ class BaseAdornmentNodeDisplay(BaseNodeDisplay[_BaseAdornmentNodeType], Generic[
                 "Unable to serialize standalone adornment nodes. Please use adornment nodes as a decorator."
             )
 
-        wrapped_node_display_class = get_node_display_class(wrapped_node)
-        wrapped_node_display = wrapped_node_display_class()
+        # Use existing node display instance from display_context if available to preserve build() data
+        wrapped_node_display = display_context.node_displays.get(wrapped_node)
+        if not wrapped_node_display:
+            wrapped_node_display_class = get_node_display_class(wrapped_node)
+            wrapped_node_display = wrapped_node_display_class()
         additional_kwargs = get_additional_kwargs(wrapped_node_display.node_id) if get_additional_kwargs else {}
         serialized_wrapped_node = wrapped_node_display.serialize(display_context, **kwargs, **additional_kwargs)
 

vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -119,5 +119,5 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
                 "output_id": str(self.output_id) if self.output_id else str(output_display.id),
                 "log_output_id": str(self.log_output_id) if self.log_output_id else str(log_output_display.id),
             },
-            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
+            **self.serialize_generic_fields(display_context),
         }

vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -2,10 +2,11 @@ from uuid import UUID
 from typing import TYPE_CHECKING, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 from vellum import FunctionDefinition, PromptBlock, RichTextChildBlock, VellumVariable
+from vellum.workflows import MCPServer
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.types.core import JsonObject
-from vellum.workflows.types.definition import DeploymentDefinition, VellumIntegrationToolDefinition
+from vellum.workflows.types.definition import DeploymentDefinition, MCPToolDefinition, VellumIntegrationToolDefinition
 from vellum.workflows.types.generics import is_workflow_class
 from vellum.workflows.utils.functions import (
     compile_function_definition,
@@ -82,6 +83,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             [
                 self._generate_function_tools(function, i, display_context)
                 for i, function in enumerate(function_definitions)
+                if not isinstance(function, (MCPServer, MCPToolDefinition))  # we don't need to serialize MCP types
             ]
             if isinstance(function_definitions, list)
             else []
@@ -147,7 +149,13 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             )
             vellum_variable_type = infer_vellum_variable_type(variable_value)
             node_inputs.append(node_input)
-            prompt_inputs.append(VellumVariable(id=str(node_input.id), key=variable_name, type=vellum_variable_type))
+            prompt_inputs.append(
+                VellumVariable(
+                    id=str(node_input.id),
+                    key=variable_name,
+                    type=vellum_variable_type,
+                )
+            )
 
         return node_inputs, prompt_inputs
 

vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py

@@ -14,7 +14,7 @@ from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.exceptions import NodeValidationError
-from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
+from vellum_ee.workflows.display.utils.vellum import compile_descriptor_annotation, infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -111,20 +111,23 @@ class BaseInlineSubworkflowNodeDisplay(
             for variable_name, variable_value in subworkflow_entries
         ]
         node_inputs_by_key = {node_input.key: node_input for node_input in node_inputs}
-        workflow_inputs = [
-            VellumVariable(
-                id=node_inputs_by_key[descriptor.name].id,
-                key=descriptor.name,
-                type=infer_vellum_variable_type(descriptor),
-                required=descriptor.instance is undefined,
-                default=(
-                    primitive_to_vellum_value(descriptor.instance).dict()
-                    if descriptor.instance is not undefined
-                    else None
-                ),
+        workflow_inputs = []
+        for descriptor in subworkflow_inputs_class:
+            schema = compile_descriptor_annotation(descriptor)
+            workflow_inputs.append(
+                VellumVariable(
+                    id=node_inputs_by_key[descriptor.name].id,
+                    key=descriptor.name,
+                    type=infer_vellum_variable_type(descriptor),
+                    required=descriptor.instance is undefined,
+                    default=(
+                        primitive_to_vellum_value(descriptor.instance).dict()
+                        if descriptor.instance is not undefined
+                        else None
+                    ),
+                    schema=schema,
+                )
             )
-            for descriptor in subworkflow_inputs_class
-        ]
 
         return node_inputs, workflow_inputs
 

vellum_ee/workflows/display/nodes/vellum/map_node.py

@@ -17,6 +17,7 @@ _MapNodeType = TypeVar("_MapNodeType", bound=MapNode)
 
 class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNodeType]):
     __serializable_inputs__ = {MapNode.items}  # type: ignore[misc]
+    __unserializable_attributes__ = {MapNode.subworkflow, MapNode.max_concurrency}  # type: ignore[misc]
 
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
@@ -89,6 +90,7 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
                 "item_input_id": item_workflow_input_id,
                 "index_input_id": index_workflow_input_id,
             },
+            "attributes": self._serialize_attributes(display_context),
             **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
 

vellum_ee/workflows/display/nodes/vellum/note_node.py

@@ -1,20 +1,35 @@
-from typing import Any, ClassVar, Dict, Generic, TypeVar, Union
+import warnings
+from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, Union, cast
 
 from vellum.workflows.nodes import NoteNode
 from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _NoteNodeType = TypeVar("_NoteNodeType", bound=NoteNode)
 
 
 class BaseNoteNodeDisplay(BaseNodeDisplay[_NoteNodeType], Generic[_NoteNodeType]):
+    # Deprecated: Define text and style on the node class instead. Will be removed in v2.0.0.
     text: ClassVar[str] = ""
     style: ClassVar[Union[Dict[str, Any], None]] = None
 
     def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
         del kwargs  # Unused parameters
         node_id = self.node_id
+        node = self._node
+
+        text = raise_if_descriptor(node.text) or ""
+        style = cast(Optional[Dict[str, Any]], raise_if_descriptor(node.style)) or None
+
+        if "text" in self.__class__.__dict__ or "style" in self.__class__.__dict__:
+            warnings.warn(
+                "Defining 'text' and 'style' on the display class is deprecated. "
+                "Define them on the node class instead. Will be removed in v2.0.0.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
         return {
             "id": str(node_id),
@@ -22,8 +37,8 @@ class BaseNoteNodeDisplay(BaseNodeDisplay[_NoteNodeType], Generic[_NoteNodeType]
             "inputs": [],
             "data": {
                 "label": self.label,
-                "text": self.text,
-                "style": self.style,
+                "text": text,
+                "style": style,
             },
             **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }

vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py

@@ -1,11 +1,13 @@
 from uuid import UUID
 from typing import Generic, Optional, TypeVar
 
+from vellum.client import Vellum as VellumClient
+from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import SubworkflowDeploymentNode
 from vellum.workflows.types.core import JsonObject
-from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -18,6 +20,37 @@ class BaseSubworkflowDeploymentNodeDisplay(
 ):
     __serializable_inputs__ = {SubworkflowDeploymentNode.subworkflow_inputs}
 
+    _deployment_id: Optional[str] = None
+    _release_tag: Optional[str] = None
+
+    def build(self, client: VellumClient) -> None:
+        node = self._node
+        deployment_descriptor_id = str(raise_if_descriptor(node.deployment))
+        self._release_tag = raise_if_descriptor(node.release_tag)
+
+        deployment_release = client.workflow_deployments.retrieve_workflow_deployment_release(
+            id=deployment_descriptor_id,
+            release_id_or_release_tag=self._release_tag,
+        )
+        self._deployment_id = str(deployment_release.deployment.id)
+        output_variables_by_key = {var.key: var for var in deployment_release.workflow_version.output_variables}
+
+        for output in node.Outputs:
+            original_output_display = self.output_display.get(output)
+            if original_output_display and not original_output_display._is_implicit:
+                continue
+
+            output_variable = output_variables_by_key.get(output.name)
+            if not output_variable or not is_valid_uuid(output_variable.id):
+                continue
+
+            output_id = UUID(output_variable.id)
+            self.output_display[output] = NodeOutputDisplay(
+                id=output_id,
+                name=output.name,
+            )
+            self._node.__output_ids__[output.name] = output_id
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
@@ -45,16 +78,6 @@ class BaseSubworkflowDeploymentNodeDisplay(
             for variable_name, variable_value in input_items
         ]
 
-        deployment_descriptor_id = str(raise_if_descriptor(node.deployment))
-        try:
-            deployment = display_context.client.workflow_deployments.retrieve(
-                id=deployment_descriptor_id,
-            )
-            deployment_id = str(deployment.id)
-        except Exception as e:
-            display_context.add_error(e)
-            deployment_id = str(uuid4_from_hash(deployment_descriptor_id))
-
         return {
             "id": str(node_id),
             "type": "SUBWORKFLOW",
@@ -65,8 +88,8 @@ class BaseSubworkflowDeploymentNodeDisplay(
                 "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
                 "target_handle_id": str(self.get_target_handle_id()),
                 "variant": "DEPLOYMENT",
-                "workflow_deployment_id": deployment_id,
-                "release_tag": raise_if_descriptor(node.release_tag),
+                "workflow_deployment_id": self._deployment_id,
+                "release_tag": self._release_tag,
             },
-            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
+            **self.serialize_generic_fields(display_context),
         }

vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py

@@ -6,6 +6,7 @@ from vellum.client.core.api_error import ApiError
 from vellum.workflows.environment import EnvironmentVariables
 from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
 from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.state.base import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.vellum.code_execution_node import BaseCodeExecutionNodeDisplay
 from vellum_ee.workflows.display.utils.exceptions import NodeValidationError
@@ -217,10 +218,13 @@ def test_serialize_node__with_environment_variable_references():
     assert other_config_input["value"]["combinator"] == "OR"
     assert len(other_config_input["value"]["rules"]) == 1
     assert other_config_input["value"]["rules"][0]["type"] == "CONSTANT_VALUE"
-    # The nested dict should contain the environment variable reference
     nested_data = other_config_input["value"]["rules"][0]["data"]
     assert nested_data["type"] == "JSON"
-    assert "nested_key" in nested_data["value"]
+    assert nested_data["value"]["type"] == "DICTIONARY_REFERENCE"
+    assert len(nested_data["value"]["entries"]) == 1
+    assert nested_data["value"]["entries"][0]["key"] == "nested_key"
+    assert nested_data["value"]["entries"][0]["value"]["type"] == "ENVIRONMENT_VARIABLE"
+    assert nested_data["value"]["entries"][0]["value"]["environment_variable"] == "NESTED_KEY"
 
 
 def test_serialize_node__with_non_exist_code_input_path_with_dry_run():
@@ -303,6 +307,22 @@ def test_serialize_node__with_non_exist_code_input_path_with_dry_run():
                     },
                     "ports": [{"id": "dc9edb2e-4392-4a2c-ab92-cc1b9c0cbd53", "name": "default", "type": "DEFAULT"}],
                     "trigger": {"id": "66e7ef63-518b-40e7-911a-e38e8bcaec81", "merge_behavior": "AWAIT_ANY"},
+                    "outputs": [
+                        {
+                            "id": "98cae9b9-45cc-4897-a0f5-df250b56c00d",
+                            "name": "result",
+                            "schema": {"type": "string"},
+                            "type": "STRING",
+                            "value": None,
+                        },
+                        {
+                            "id": "66c06c97-a9d1-4abf-840f-3f6c29709612",
+                            "name": "log",
+                            "schema": {"type": "string"},
+                            "type": "STRING",
+                            "value": None,
+                        },
+                    ],
                 },
             ],
             "edges": [
@@ -326,3 +346,43 @@ def test_serialize_node__with_non_exist_code_input_path_with_dry_run():
         "state_variables": [],
         "output_variables": [],
     }
+
+
+def test_serialize_node__with_custom_output_type():
+    # GIVEN a code node with a custom output type
+    class MyNode(CodeExecutionNode[BaseState, dict[str, int]]):
+        code = """\
+return {
+    "hello": 1,
+}
+"""
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node's outputs should serialize correctly
+    my_code_execution_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "CODE_EXECUTION"
+    )
+    assert my_code_execution_node["outputs"] == [
+        {
+            "id": "01de8e8b-5e0e-4344-93b0-e002bbaed840",
+            "name": "result",
+            "value": None,
+            "type": "JSON",
+            "schema": {"type": "object", "additionalProperties": {"type": "integer"}},
+        },
+        {
+            "id": "64bf62e0-adc7-48cc-b689-8bf9b7e4eeef",
+            "name": "log",
+            "value": None,
+            "type": "STRING",
+            "schema": {"type": "string"},
+        },
+    ]
+    assert my_code_execution_node["data"]["output_type"] == "JSON"