vellum-ai 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

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

@@ -1,11 +1,17 @@
 from uuid import UUID
-from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 from vellum import FunctionDefinition, PromptBlock, RichTextChildBlock, VellumVariable
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.types.core import JsonObject
-from vellum.workflows.utils.functions import compile_function_definition
+from vellum.workflows.types.definition import DeploymentDefinition
+from vellum.workflows.types.generics import is_workflow_class
+from vellum.workflows.utils.functions import (
+    compile_function_definition,
+    compile_inline_workflow_function_definition,
+    compile_workflow_deployment_function_definition,
+)
 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.utils import raise_if_descriptor
@@ -14,6 +20,9 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
 
 def _contains_descriptors(obj):
     """Check if an object contains any descriptors or references that need special handling."""
@@ -69,7 +78,10 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             ]
 
         functions = (
-            [self._generate_function_tools(function, i) for i, function in enumerate(function_definitions)]
+            [
+                self._generate_function_tools(function, i, display_context)
+                for i, function in enumerate(function_definitions)
+            ]
             if isinstance(function_definitions, list)
             else []
         )
@@ -145,10 +157,24 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
 
         return node_inputs, prompt_inputs
 
-    def _generate_function_tools(self, function: Union[FunctionDefinition, Callable], index: int) -> JsonObject:
-        normalized_functions = (
-            function if isinstance(function, FunctionDefinition) else compile_function_definition(function)
-        )
+    def _generate_function_tools(
+        self,
+        function: Union[FunctionDefinition, Callable, DeploymentDefinition, Type["BaseWorkflow"]],
+        index: int,
+        display_context: WorkflowDisplayContext,
+    ) -> JsonObject:
+        if isinstance(function, FunctionDefinition):
+            normalized_functions = function
+        elif is_workflow_class(function):
+            normalized_functions = compile_inline_workflow_function_definition(function)
+        elif callable(function):
+            normalized_functions = compile_function_definition(function)
+        elif isinstance(function, DeploymentDefinition):
+            normalized_functions = compile_workflow_deployment_function_definition(
+                function.model_dump(), display_context.client
+            )
+        else:
+            raise ValueError(f"Unsupported function type: {type(function)}")
         return {
             "id": str(uuid4_from_hash(f"{self.node_id}-FUNCTION_DEFINITION-{index}")),
             "block_type": "FUNCTION_DEFINITION",

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

@@ -1,13 +1,31 @@
+from datetime import datetime
+
 from vellum.client.types.prompt_parameters import PromptParameters
+from vellum.client.types.release_review_reviewer import ReleaseReviewReviewer
+from vellum.client.types.workflow_deployment_release import (
+    ReleaseEnvironment,
+    ReleaseReleaseTag,
+    SlimReleaseReview,
+    WorkflowDeploymentRelease,
+    WorkflowDeploymentReleaseWorkflowDeployment,
+    WorkflowDeploymentReleaseWorkflowVersion,
+)
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
 from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
 from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_router_node, create_tool_prompt_node
+from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.types.definition import AuthorizationType, EnvironmentVariableReference, MCPServer
+from vellum.workflows.types.definition import (
+    AuthorizationType,
+    DeploymentDefinition,
+    EnvironmentVariableReference,
+    MCPServer,
+)
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -185,6 +203,58 @@ def test_serialize_node__tool_calling_node__mcp_server_api_key():
     }
 
 
+def test_serialize_node__tool_calling_node__mcp_server_no_authorization():
+    # GIVEN a tool calling node with an mcp server
+    class MyToolCallingNode(ToolCallingNode):
+        functions = [
+            MCPServer(
+                name="my-mcp-server",
+                url="https://my-mcp-server.com",
+            )
+        ]
+
+    # AND a workflow with the tool calling node
+    class Workflow(BaseWorkflow):
+        graph = MyToolCallingNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the mcp server
+    my_tool_calling_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(MyToolCallingNode.__id__)
+    )
+
+    functions_attribute = next(
+        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
+    )
+
+    assert functions_attribute == {
+        "id": "c8957551-cb3d-49af-8053-acd256c1d852",
+        "name": "functions",
+        "value": {
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": [
+                    {
+                        "type": "MCP_SERVER",
+                        "name": "my-mcp-server",
+                        "url": "https://my-mcp-server.com",
+                        "authorization_type": None,
+                        "bearer_token_value": None,
+                        "api_key_header_key": None,
+                        "api_key_header_value": None,
+                    }
+                ],
+            },
+        },
+    }
+
+
 def test_serialize_tool_router_node():
     """
     Test that the tool router node created by create_router_node serializes successfully.
@@ -406,3 +476,171 @@ def test_serialize_node__tool_calling_node__subworkflow_with_parent_input_refere
             "combinator": "OR",
         },
     }
+
+
+def test_serialize_tool_prompt_node_with_inline_workflow():
+    """
+    Test that the tool prompt node created by create_tool_prompt_node serializes successfully with inline workflow.
+    """
+
+    # GIVEN a simple inline workflow for tool calling
+    class SimpleWorkflowInputs(BaseInputs):
+        message: str
+
+    class SimpleNode(BaseNode):
+        message = SimpleWorkflowInputs.message
+
+        class Outputs(BaseOutputs):
+            result: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(result=f"Processed: {self.message}")
+
+    class SimpleInlineWorkflow(BaseWorkflow[SimpleWorkflowInputs, BaseState]):
+        """A simple workflow for testing inline tool serialization."""
+
+        graph = SimpleNode
+
+        class Outputs(BaseOutputs):
+            result = SimpleNode.Outputs.result
+
+    # WHEN we create a tool prompt node using create_tool_prompt_node with inline workflow
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[SimpleInlineWorkflow],
+        prompt_inputs=None,
+        parameters=PromptParameters(),
+    )
+
+    tool_prompt_node_display_class = get_node_display_class(tool_prompt_node)
+    tool_prompt_node_display = tool_prompt_node_display_class()
+
+    # AND we create a workflow that uses this tool prompt node
+    class TestWorkflow(BaseWorkflow[BaseInputs, ToolCallingState]):
+        graph = tool_prompt_node
+
+    # WHEN we serialize the entire workflow
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+    display_context = workflow_display.display_context
+    serialized_tool_prompt_node = tool_prompt_node_display.serialize(display_context)
+
+    # THEN prompt inputs should be serialized correctly
+    attributes = serialized_tool_prompt_node["attributes"]
+    assert isinstance(attributes, list)
+    prompt_inputs_attr = next(
+        (attr for attr in attributes if isinstance(attr, dict) and attr["name"] == "prompt_inputs"), None
+    )
+    assert prompt_inputs_attr == {
+        "id": "bc1320a2-23e4-4238-8b00-efbf88e91856",
+        "name": "prompt_inputs",
+        "value": {
+            "type": "DICTIONARY_REFERENCE",
+            "entries": [
+                {
+                    "id": "76ceec7b-ec37-474f-ba38-2bfd27cecc5d",
+                    "key": "chat_history",
+                    "value": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": []}},
+                        "operator": "concat",
+                        "rhs": {
+                            "type": "WORKFLOW_STATE",
+                            "state_variable_id": "7a1caaf5-99df-487a-8b2d-6512df2d871a",
+                        },
+                    },
+                }
+            ],
+        },
+    }
+
+
+def test_serialize_tool_prompt_node_with_workflow_deployment(vellum_client):
+    """
+    Test that the tool prompt node serializes successfully with a workflow deployment.
+    """
+    vellum_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = WorkflowDeploymentRelease(
+        id="test-id",
+        created=datetime.now(),
+        environment=ReleaseEnvironment(
+            id="test-id",
+            name="test-name",
+            label="test-label",
+        ),
+        created_by=None,
+        workflow_version=WorkflowDeploymentReleaseWorkflowVersion(
+            id="test-id",
+            input_variables=[],
+            output_variables=[],
+        ),
+        deployment=WorkflowDeploymentReleaseWorkflowDeployment(name="test-name"),
+        description="test-description",
+        release_tags=[
+            ReleaseReleaseTag(
+                name="test-name",
+                source="USER",
+            )
+        ],
+        reviews=[
+            SlimReleaseReview(
+                id="test-id",
+                created=datetime.now(),
+                reviewer=ReleaseReviewReviewer(
+                    id="test-id",
+                    full_name="test-name",
+                ),
+                state="APPROVED",
+            )
+        ],
+    )
+
+    # GIVEN a workflow deployment
+    workflow_deployment = DeploymentDefinition(
+        deployment="test-deployment",
+        release_tag="test-release-tag",
+    )
+
+    # WHEN we create a tool prompt node using create_tool_prompt_node with a workflow deployment
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[workflow_deployment],
+        prompt_inputs=None,
+        parameters=PromptParameters(),
+    )
+
+    tool_prompt_node_display_class = get_node_display_class(tool_prompt_node)
+    tool_prompt_node_display = tool_prompt_node_display_class()
+
+    # AND we create a workflow that uses this tool prompt node
+    class TestWorkflow(BaseWorkflow[BaseInputs, ToolCallingState]):
+        graph = tool_prompt_node
+
+    # WHEN we serialize the entire workflow
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+    display_context = workflow_display.display_context
+    serialized_tool_prompt_node = tool_prompt_node_display.serialize(display_context)
+
+    # THEN functions attribute should be serialized correctly
+    attributes = serialized_tool_prompt_node["attributes"]
+    assert isinstance(attributes, list)
+    functions_attr = next((attr for attr in attributes if isinstance(attr, dict) and attr["name"] == "functions"), None)
+    assert functions_attr == {
+        "id": "6326ccc4-7cf6-4235-ba3c-a6e860b0c48b",
+        "name": "functions",
+        "value": {
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": [
+                    {
+                        "type": "WORKFLOW_DEPLOYMENT",
+                        "name": "test-name",
+                        "description": "test-description",
+                        "deployment": "test-deployment",
+                        "release_tag": "test-release-tag",
+                    }
+                ],
+            },
+        },
+    }

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -2,7 +2,8 @@ from uuid import UUID
 from typing import Dict
 
 from vellum.workflows.inputs import BaseInputs
-from vellum.workflows.nodes import BaseNode
+from vellum.workflows.nodes import BaseNode, InlineSubworkflowNode
+from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state import BaseState
@@ -327,3 +328,54 @@ def test_serialize__port_with_lazy_reference():
             },
         }
     ]
+
+
+def test_global_propagation_deep_nested_subworkflows():
+    # GIVEN the root workflow, a middle workflow, and an inner workflow
+
+    class RootInputs(BaseInputs):
+        root_param: str
+
+    class MiddleInputs(BaseInputs):
+        middle_param: str
+
+    class InnerInputs(BaseInputs):
+        inner_param: str
+
+    class InnerNode(BaseNode):
+        class Outputs(BaseOutputs):
+            done: bool
+
+        def run(self) -> Outputs:
+            return self.Outputs(done=True)
+
+    class InnerWorkflow(BaseWorkflow[InnerInputs, BaseState]):
+        graph = InnerNode
+
+    class MiddleInlineSubworkflowNode(InlineSubworkflowNode):
+        subworkflow_inputs = {"inner_param": "x"}
+        subworkflow = InnerWorkflow
+
+    class MiddleWorkflow(BaseWorkflow[MiddleInputs, BaseState]):
+        graph = MiddleInlineSubworkflowNode
+
+    class OuterInlineSubworkflowNode(InlineSubworkflowNode):
+        subworkflow_inputs = {"middle_param": "y"}
+        subworkflow = MiddleWorkflow
+
+    class RootWorkflow(BaseWorkflow[RootInputs, BaseState]):
+        graph = OuterInlineSubworkflowNode
+
+    # WHEN we build the displays
+    root_display = get_workflow_display(workflow_class=RootWorkflow)
+    middle_display = get_workflow_display(
+        workflow_class=MiddleWorkflow, parent_display_context=root_display.display_context
+    )
+    inner_display = get_workflow_display(
+        workflow_class=InnerWorkflow, parent_display_context=middle_display.display_context
+    )
+
+    # THEN the deepest display must include root + middle + inner inputs in its GLOBAL view
+    inner_global_names = {ref.name for ref in inner_display.display_context.global_workflow_input_displays.keys()}
+
+    assert inner_global_names == {"middle_param", "inner_param", "root_param"}

vellum_ee/workflows/display/utils/expressions.py

@@ -11,6 +11,7 @@ from vellum.workflows.expressions.and_ import AndExpression
 from vellum.workflows.expressions.begins_with import BeginsWithExpression
 from vellum.workflows.expressions.between import BetweenExpression
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.expressions.concat import ConcatExpression
 from vellum.workflows.expressions.contains import ContainsExpression
 from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
 from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
@@ -105,6 +106,8 @@ def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperato
         return "+"
     elif isinstance(descriptor, MinusExpression):
         return "-"
+    elif isinstance(descriptor, ConcatExpression):
+        return "concat"
     else:
         raise ValueError(f"Unsupported descriptor type: {descriptor}")
 
@@ -171,6 +174,7 @@ def _serialize_condition(display_context: "WorkflowDisplayContext", condition: B
             AndExpression,
             BeginsWithExpression,
             CoalesceExpression,
+            ConcatExpression,
             ContainsExpression,
             DoesNotBeginWithExpression,
             DoesNotContainExpression,

vellum_ee/workflows/display/utils/registry.py

@@ -1,10 +1,14 @@
+from uuid import UUID
 from typing import TYPE_CHECKING, Dict, Optional, Type
 
+from vellum.workflows.events.types import BaseEvent
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.workflows.base import BaseWorkflow
 
 if TYPE_CHECKING:
+    from vellum.workflows.events.types import ParentContext
     from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+    from vellum_ee.workflows.display.types import WorkflowDisplayContext
     from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
@@ -14,6 +18,9 @@ _workflow_display_registry: Dict[Type[BaseWorkflow], Type["BaseWorkflowDisplay"]
 # Used to store the mapping between node types and their display classes
 _node_display_registry: Dict[Type[BaseNode], Type["BaseNodeDisplay"]] = {}
 
+# Registry to store active workflow display contexts by span ID for nested workflow inheritance
+_active_workflow_display_contexts: Dict[UUID, "WorkflowDisplayContext"] = {}
+
 
 def get_from_workflow_display_registry(workflow_class: Type[BaseWorkflow]) -> Optional[Type["BaseWorkflowDisplay"]]:
     return _workflow_display_registry.get(workflow_class)
@@ -35,3 +42,42 @@ def get_from_node_display_registry(node_class: Type[BaseNode]) -> Optional[Type[
 
 def register_node_display_class(node_class: Type[BaseNode], node_display_class: Type["BaseNodeDisplay"]) -> None:
     _node_display_registry[node_class] = node_display_class
+
+
+def register_workflow_display_context(span_id: UUID, display_context: "WorkflowDisplayContext") -> None:
+    """Register a workflow display context by span ID for nested workflow inheritance."""
+    _active_workflow_display_contexts[span_id] = display_context
+
+
+def _get_parent_display_context_for_span(span_id: UUID) -> Optional["WorkflowDisplayContext"]:
+    """Get the parent display context for a given span ID."""
+    return _active_workflow_display_contexts.get(span_id)
+
+
+def get_parent_display_context_from_event(event: BaseEvent) -> Optional["WorkflowDisplayContext"]:
+    """Extract parent display context from an event by traversing the parent chain.
+
+    This function traverses up the parent chain starting from the event's parent,
+    looking for workflow parents and attempting to get their display context.
+
+    Args:
+        event: The event to extract parent display context from
+
+    Returns:
+        The parent workflow display context if found, None otherwise
+    """
+    if not event.parent:
+        return None
+
+    current_parent: Optional["ParentContext"] = event.parent
+    while current_parent:
+        if current_parent.type == "WORKFLOW":
+            # Found a parent workflow, try to get its display context
+            parent_span_id = current_parent.span_id
+            parent_display_context = _get_parent_display_context_for_span(parent_span_id)
+            if parent_display_context:
+                return parent_display_context
+        # Move up the parent chain
+        current_parent = current_parent.parent
+
+    return None

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -528,7 +528,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         workflow_input_displays: WorkflowInputsDisplays = {}
         # If we're dealing with a nested workflow, then it should have access to the inputs of its parents.
         global_workflow_input_displays = (
-            copy(self._parent_display_context.workflow_input_displays) if self._parent_display_context else {}
+            copy(self._parent_display_context.global_workflow_input_displays) if self._parent_display_context else {}
         )
         for workflow_input in self._workflow.get_inputs_class():
             workflow_input_display_overrides = self.inputs_display.get(workflow_input)

vellum_ee/workflows/tests/test_registry.py

@@ -0,0 +1,169 @@
+from datetime import datetime, timezone
+from uuid import uuid4
+
+from vellum.workflows.events.types import NodeParentContext, WorkflowParentContext
+from vellum.workflows.events.workflow import WorkflowExecutionInitiatedBody, WorkflowExecutionInitiatedEvent
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.utils.registry import (
+    get_parent_display_context_from_event,
+    register_workflow_display_context,
+)
+
+
+class MockInputs(BaseInputs):
+    pass
+
+
+class MockState(BaseState):
+    pass
+
+
+class MockNode(BaseNode):
+    pass
+
+
+class MockWorkflow(BaseWorkflow[MockInputs, MockState]):
+    pass
+
+
+class MockWorkflowDisplayContext:
+    pass
+
+
+def test_get_parent_display_context_from_event__no_parent():
+    """Test event with no parent returns None"""
+    # GIVEN a workflow execution initiated event with no parent
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=None,  # No parent
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should return None
+    assert result is None
+
+
+def test_get_parent_display_context_from_event__non_workflow_parent():
+    """Test event with non-workflow parent continues traversal"""
+    # GIVEN an event with a non-workflow parent (NodeParentContext)
+    non_workflow_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=None)
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=non_workflow_parent,
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should return None (no workflow parent found)
+    assert result is None
+
+
+def test_get_parent_display_context_from_event__nested_workflow_parents():
+    """Test event with nested workflow parents traverses correctly"""
+    # GIVEN a chain of nested contexts:
+    # Event -> WorkflowParent -> NodeParent -> MiddleWorkflowParent -> NodeParent
+
+    # Top level workflow parent
+    top_workflow_span_id = uuid4()
+    top_context = MockWorkflowDisplayContext()
+    register_workflow_display_context(top_workflow_span_id, top_context)  # type: ignore[arg-type]
+
+    top_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=top_workflow_span_id, parent=None
+    )
+
+    top_node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=top_workflow_parent)
+
+    # AND middle workflow parent (no display context)
+    middle_workflow_span_id = uuid4()
+    middle_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=middle_workflow_span_id, parent=top_node_parent
+    )
+
+    # AND node parent between middle workflow and event
+    node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=middle_workflow_parent)
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=node_parent,
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should find the top-level workflow context
+    assert result == top_context
+
+
+def test_get_parent_display_context_from_event__middle_workflow_has_context():
+    """Test event returns middle workflow context when it's the first one with registered context"""
+    # GIVEN a chain of nested contexts:
+    # Event -> WorkflowParent -> NodeParent -> MiddleWorkflowParent -> NodeParent
+
+    top_workflow_span_id = uuid4()
+    top_context = MockWorkflowDisplayContext()
+    register_workflow_display_context(top_workflow_span_id, top_context)  # type: ignore[arg-type]
+
+    top_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=top_workflow_span_id, parent=None
+    )
+
+    # AND node parent between top workflow and middle workflow
+    top_node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=top_workflow_parent)
+
+    # AND middle workflow parent
+    middle_workflow_span_id = uuid4()
+    middle_context = MockWorkflowDisplayContext()
+    register_workflow_display_context(middle_workflow_span_id, middle_context)  # type: ignore[arg-type]
+
+    middle_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=middle_workflow_span_id, parent=top_node_parent
+    )
+
+    node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=middle_workflow_parent)
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=node_parent,
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should find the MIDDLE workflow context
+    assert result == middle_context