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

vellum/workflows/runner/runner.py

@@ -1,15 +1,35 @@
 from collections import defaultdict
+from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass
 import logging
 from queue import Empty, Queue
+import sys
 from threading import Event as ThreadingEvent, Thread
+import traceback
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generator,
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
+from vellum.client.core.api_error import ApiError
 from vellum.workflows.constants import undefined
 from vellum.workflows.context import ExecutionContext, execution_context, get_execution_context
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.exceptions import InvalidExpressionException
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
 from vellum.workflows.events import (
     NodeExecutionFulfilledEvent,
@@ -17,20 +37,29 @@ from vellum.workflows.events import (
     NodeExecutionRejectedEvent,
     NodeExecutionStreamingEvent,
     WorkflowEvent,
-    WorkflowEventStream,
+    WorkflowEventGenerator,
     WorkflowExecutionFulfilledEvent,
     WorkflowExecutionInitiatedEvent,
     WorkflowExecutionRejectedEvent,
     WorkflowExecutionStreamingEvent,
 )
 from vellum.workflows.events.node import (
+    NodeEvent,
     NodeExecutionFulfilledBody,
     NodeExecutionInitiatedBody,
     NodeExecutionRejectedBody,
     NodeExecutionStreamingBody,
 )
-from vellum.workflows.events.types import BaseEvent, NodeParentContext, ParentContext, WorkflowParentContext
+from vellum.workflows.events.types import (
+    BaseEvent,
+    NodeParentContext,
+    ParentContext,
+    SpanLink,
+    WorkflowParentContext,
+    default_serializer,
+)
 from vellum.workflows.events.workflow import (
+    WorkflowEventStream,
     WorkflowExecutionFulfilledBody,
     WorkflowExecutionInitiatedBody,
     WorkflowExecutionPausedBody,
@@ -43,6 +72,7 @@ from vellum.workflows.events.workflow import (
     WorkflowExecutionStreamingBody,
 )
 from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
+from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import NodeRunResponse
 from vellum.workflows.nodes.mocks import MockNodeExecutionArg
@@ -50,7 +80,13 @@ from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references import ExternalInputReference, OutputReference
+from vellum.workflows.references.state_value import StateValueReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.delta import StateDelta
+from vellum.workflows.triggers.base import BaseTrigger
+from vellum.workflows.triggers.integration import IntegrationTrigger
+from vellum.workflows.triggers.manual import ManualTrigger
+from vellum.workflows.types.core import CancelSignal
 from vellum.workflows.types.generics import InputsType, OutputsType, StateType
 
 if TYPE_CHECKING:
@@ -58,7 +94,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-RunFromNodeArg = Sequence[Type[BaseNode]]
+RunFromNodeArg = Sequence[Union[Type[BaseNode], UUID]]
 ExternalInputsArg = Dict[ExternalInputReference, Any]
 BackgroundThreadItem = Union[BaseState, WorkflowEvent, None]
 
@@ -79,10 +115,15 @@ class WorkflowRunner(Generic[StateType]):
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
-        cancel_signal: Optional[ThreadingEvent] = None,
+        previous_execution_id: Optional[Union[str, UUID]] = None,
+        cancel_signal: Optional[CancelSignal] = None,
         node_output_mocks: Optional[MockNodeExecutionArg] = None,
         max_concurrency: Optional[int] = None,
+        timeout: Optional[float] = None,
         init_execution_context: Optional[ExecutionContext] = None,
+        trigger: Optional[BaseTrigger] = None,
+        execution_id: Optional[UUID] = None,
+        event_max_size: Optional[int] = None,
     ):
         if state and external_inputs:
             raise ValueError("Can only run a Workflow providing one of state or external inputs, not both")
@@ -90,22 +131,35 @@ class WorkflowRunner(Generic[StateType]):
         self.workflow = workflow
         self._is_resuming = False
         self._should_emit_initial_state = True
+        self._span_link_info: Optional[Tuple[str, str, str, str]] = None
         if entrypoint_nodes:
-            if len(list(entrypoint_nodes)) > 1:
-                raise ValueError("Cannot resume from multiple nodes")
+            nodes_by_id = {node.__id__: node for node in self.workflow.get_all_nodes()}
+
+            resolved_nodes = []
+            for item in entrypoint_nodes:
+                if isinstance(item, UUID):
+                    matching_node = nodes_by_id.get(item)
+                    if matching_node is None:
+                        raise ValueError(f"No node found with UUID {item}")
+                    resolved_nodes.append(matching_node)
+                else:
+                    resolved_nodes.append(item)
+
+            if len(list(resolved_nodes)) > 1:
+                raise WorkflowInitializationException("Cannot resume from multiple nodes")
 
             # TODO: Support resuming from multiple nodes
             # https://app.shortcut.com/vellum/story/4408
-            node = next(iter(entrypoint_nodes))
+            node = next(iter(resolved_nodes))
             if state:
                 self._initial_state = deepcopy(state)
-                self._initial_state.meta.span_id = uuid4()
+                self._initial_state.meta.span_id = execution_id or uuid4()
                 self._initial_state.meta.workflow_definition = self.workflow.__class__
             else:
                 self._initial_state = self.workflow.get_state_at_node(node)
-            self._entrypoints = entrypoint_nodes
+            self._entrypoints = resolved_nodes
         elif external_inputs:
-            self._initial_state = self.workflow.get_most_recent_state()
+            self._initial_state = self.workflow.get_most_recent_state(execution_id)
             for descriptor, value in external_inputs.items():
                 if not any(isinstance(value, type_) for type_ in descriptor.types):
                     raise NodeException(
@@ -120,21 +174,89 @@ class WorkflowRunner(Generic[StateType]):
                 if issubclass(ei.inputs_class.__parent_class__, BaseNode)
             ]
             self._is_resuming = True
+        elif previous_execution_id:
+            if not self.workflow.resolvers:
+                raise WorkflowInitializationException(
+                    message=f"No resolvers configured to load initial state for execution ID: {previous_execution_id}",
+                    workflow_definition=self.workflow.__class__,
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                )
+
+            resolver_failed = True
+            for resolver in self.workflow.resolvers:
+                try:
+                    load_state_result = resolver.load_state(previous_execution_id)
+
+                    if load_state_result is not None:
+                        if execution_id:
+                            load_state_result.state.meta.span_id = execution_id
+
+                        state_class = self.workflow.get_state_class()
+                        if isinstance(load_state_result.state, state_class):
+                            self._initial_state = load_state_result.state
+                            normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
+                            self._initial_state.meta.workflow_inputs = normalized_inputs
+                            self._initial_state.meta.workflow_definition = self.workflow.__class__
+                            self._span_link_info = (
+                                load_state_result.previous_trace_id,
+                                load_state_result.previous_span_id,
+                                load_state_result.root_trace_id,
+                                load_state_result.root_span_id,
+                            )
+                            resolver_failed = False
+                            break
+                except Exception as e:
+                    logger.warning(f"Failed to load state from resolver {type(resolver).__name__}: {e}")
+                    continue
+
+            if resolver_failed:
+                logger.warning(
+                    f"All resolvers failed to load initial state for execution ID: {previous_execution_id}. "
+                    "Falling back to default state."
+                )
+                normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
+                self._initial_state = self.workflow.get_default_state(normalized_inputs, execution_id)
+
+            self._entrypoints = self.workflow.get_entrypoints()
+        elif trigger:
+            # When trigger is provided, set up default state and filter entrypoints by trigger type
+            default_inputs = deepcopy(inputs) if inputs else None
+            if state:
+                self._initial_state = deepcopy(state)
+                self._initial_state.meta.workflow_inputs = default_inputs
+                self._initial_state.meta.span_id = execution_id or uuid4()
+                self._initial_state.meta.workflow_definition = self.workflow.__class__
+            else:
+                self._initial_state = self.workflow.get_default_state(
+                    default_inputs,
+                    execution_id,
+                    trigger_attributes={},
+                )
+                self._should_emit_initial_state = False
+
+            # Validate and bind trigger, then filter entrypoints
+            self._validate_and_bind_trigger(trigger)
+            self._entrypoints = self.workflow.get_entrypoints()
+            self._filter_entrypoints_for_trigger(trigger)
         else:
+            # Default case: no entrypoint overrides and no trigger
             normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
             if state:
                 self._initial_state = deepcopy(state)
                 self._initial_state.meta.workflow_inputs = normalized_inputs
-                self._initial_state.meta.span_id = uuid4()
+                self._initial_state.meta.span_id = execution_id or uuid4()
                 self._initial_state.meta.workflow_definition = self.workflow.__class__
             else:
-                self._initial_state = self.workflow.get_default_state(normalized_inputs)
+                self._initial_state = self.workflow.get_default_state(normalized_inputs, execution_id)
                 # We don't want to emit the initial state on the base case of Workflow Runs, since
                 # all of that data is redundant and is derivable. It also clearly communicates that
                 # there was no initial state provided by the user to invoke the workflow.
                 self._should_emit_initial_state = False
             self._entrypoints = self.workflow.get_entrypoints()
 
+            # Check if workflow requires a trigger but none was provided
+            self._validate_no_trigger_provided()
+
         # This queue is responsible for sending events from WorkflowRunner to the outside world
         self._workflow_event_outer_queue: Queue[WorkflowEvent] = Queue()
 
@@ -153,17 +275,171 @@ class WorkflowRunner(Generic[StateType]):
 
         self._active_nodes_by_execution_id: Dict[UUID, ActiveNode[StateType]] = {}
         self._cancel_signal = cancel_signal
+        self._timeout = timeout
         self._execution_context = init_execution_context or get_execution_context()
+        self._trigger = trigger
+        self._event_max_size = event_max_size
 
         setattr(
             self._initial_state,
             "__snapshot_callback__",
-            lambda s: self._snapshot_state(s),
+            lambda s, d: self._snapshot_state(s, d),
         )
         self.workflow.context._register_event_queue(self._workflow_event_inner_queue)
         self.workflow.context._register_node_output_mocks(node_output_mocks or [])
+        self.workflow.context._register_event_max_size(event_max_size)
+
+        self._outputs_listening_to_state = [
+            descriptor for descriptor in self.workflow.Outputs if isinstance(descriptor.instance, StateValueReference)
+        ]
+
+        self._background_thread: Optional[Thread] = None
+        self._cancel_thread: Optional[Thread] = None
+        self._timeout_thread: Optional[Thread] = None
+
+    def _has_manual_trigger(self) -> bool:
+        """Check if workflow has ManualTrigger."""
+        for subgraph in self.workflow.get_subgraphs():
+            for trigger in subgraph.triggers:
+                if issubclass(trigger, ManualTrigger):
+                    return True
+        return False
+
+    def _get_entrypoints_for_trigger_type(self, trigger_class: Type) -> List[Type[BaseNode]]:
+        """Get all entrypoints connected to a specific trigger type.
+
+        Allows subclasses: if trigger_class is a subclass of any declared trigger,
+        returns those entrypoints.
+        """
+        seen: Set[Type[BaseNode]] = set()
+        entrypoints: List[Type[BaseNode]] = []
+        for subgraph in self.workflow.get_subgraphs():
+            for trigger in subgraph.triggers:
+                # Check if the provided trigger_class is a subclass of the declared trigger
+                # This allows runtime instances to be subclasses of what's declared in the workflow
+                if issubclass(trigger_class, trigger):
+                    for entrypoint in subgraph.entrypoints:
+                        if entrypoint not in seen:
+                            seen.add(entrypoint)
+                            entrypoints.append(entrypoint)
+        return entrypoints
+
+    def _validate_and_bind_trigger(self, trigger: BaseTrigger) -> None:
+        """
+        Validate that trigger is compatible with workflow and bind it to state.
+
+        Supports all trigger types derived from BaseTrigger:
+        - IntegrationTrigger instances (Slack, Gmail, etc.)
+        - ManualTrigger instances (explicit manual execution)
+        - ScheduledTrigger instances (time-based triggers)
+        - Any future trigger types
+
+        Raises:
+            WorkflowInitializationException: If trigger type is not compatible with workflow
+        """
+        trigger_class = type(trigger)
+
+        # Search for a compatible trigger type in the workflow
+        found_compatible_trigger = False
+        has_any_triggers = False
+        incompatible_trigger_names: List[str] = []
+
+        for subgraph in self.workflow.get_subgraphs():
+            for declared_trigger in subgraph.triggers:
+                has_any_triggers = True
+                # Allow subclasses: if workflow declares BaseSlackTrigger, accept SpecificSlackTrigger instances
+                if issubclass(trigger_class, declared_trigger):
+                    found_compatible_trigger = True
+                    break
+                else:
+                    incompatible_trigger_names.append(declared_trigger.__name__)
+
+            if found_compatible_trigger:
+                break
+
+        # Special case: workflows with no explicit triggers implicitly support ManualTrigger
+        if not has_any_triggers and not isinstance(trigger, ManualTrigger):
+            raise WorkflowInitializationException(
+                message=f"Provided trigger type {trigger_class.__name__} is not compatible with workflow. "
+                f"Workflow has no explicit triggers and only supports ManualTrigger.",
+                workflow_definition=self.workflow.__class__,
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+
+        # Validate that we found a compatible trigger type
+        if has_any_triggers and not found_compatible_trigger:
+            raise WorkflowInitializationException(
+                message=f"Provided trigger type {trigger_class.__name__} is not compatible with workflow triggers. "
+                f"Workflow has: {sorted(set(incompatible_trigger_names))}",
+                workflow_definition=self.workflow.__class__,
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+
+        # Bind trigger to state (works for all trigger types via BaseTrigger.bind_to_state)
+        trigger.bind_to_state(self._initial_state)
+
+    def _filter_entrypoints_for_trigger(self, trigger: BaseTrigger) -> None:
+        """
+        Filter entrypoints to those connected to the specific trigger type.
+
+        Uses the specific trigger subclass, not the parent class, allowing workflows
+        with multiple triggers to route to the correct path.
+        """
+        trigger_class = type(trigger)
+        specific_entrypoints = self._get_entrypoints_for_trigger_type(trigger_class)
+        if specific_entrypoints:
+            self._entrypoints = specific_entrypoints
+
+    def _validate_no_trigger_provided(self) -> None:
+        """
+        Validate that workflow can run without a trigger.
+
+        If workflow has IntegrationTrigger(s) but no ManualTrigger, it requires a trigger instance.
+        If workflow has both, filter entrypoints to ManualTrigger path only.
+
+        Raises:
+            WorkflowInitializationException: If workflow requires trigger but none was provided
+        """
+        # Collect all IntegrationTrigger types in the workflow
+        workflow_integration_triggers = []
+        for subgraph in self.workflow.get_subgraphs():
+            for trigger_type in subgraph.triggers:
+                if issubclass(trigger_type, IntegrationTrigger):
+                    workflow_integration_triggers.append(trigger_type)
+
+    @contextmanager
+    def _httpx_logger_with_span_id(self) -> Iterator[None]:
+        """
+        Context manager to append the current execution's span ID to httpx logger messages.
+
+        This is used when making API requests via the Vellum client to include
+        the execution's span ID in the httpx request logs for better traceability.
+        """
+        httpx_logger = logging.getLogger("httpx")
+
+        class SpanIdFilter(logging.Filter):
+            def filter(self, record: logging.LogRecord) -> bool:
+                if record.name == "httpx" and "[span_id=" not in record.msg:
+                    context = get_execution_context()
+                    if context.parent_context:
+                        span_id = str(context.parent_context.span_id)
+                        record.msg = f"{record.msg} [span_id={span_id}]"
+                return True
+
+        span_filter = SpanIdFilter()
+        httpx_logger.addFilter(span_filter)
+
+        try:
+            yield
+        finally:
+            httpx_logger.removeFilter(span_filter)
+
+    def _snapshot_state(self, state: StateType, deltas: List[StateDelta]) -> StateType:
+        execution = get_execution_context()
+        edited_by = None
+        if execution.parent_context and hasattr(execution.parent_context, "node_definition"):
+            edited_by = execution.parent_context.node_definition
 
-    def _snapshot_state(self, state: StateType) -> StateType:
         self._workflow_event_inner_queue.put(
             WorkflowExecutionSnapshottedEvent(
                 trace_id=self._execution_context.trace_id,
@@ -171,31 +447,78 @@ class WorkflowRunner(Generic[StateType]):
                 body=WorkflowExecutionSnapshottedBody(
                     workflow_definition=self.workflow.__class__,
                     state=state,
+                    edited_by=edited_by,
                 ),
                 parent=self._execution_context.parent_context,
             )
         )
+
+        delta_names = [d.name for d in deltas]
+        for descriptor in self._outputs_listening_to_state:
+            if not isinstance(descriptor.instance, StateValueReference):
+                continue
+
+            if descriptor.instance.name not in delta_names:
+                continue
+
+            resolved_delta = descriptor.instance.resolve(state)
+            if resolved_delta is undefined:
+                continue
+
+            self._workflow_event_outer_queue.put(
+                self._stream_workflow_event(
+                    BaseOutput(
+                        name=descriptor.name,
+                        # We may want to switch this to using the delta from the state delta
+                        # instead of the state value. In the short term, we need to show
+                        # the full conversation of the tool calling node's delta's.
+                        delta=resolved_delta,
+                    )
+                )
+            )
+
         self.workflow._store.append_state_snapshot(state)
         self._background_thread_queue.put(state)
         return state
 
     def _emit_event(self, event: WorkflowEvent) -> WorkflowEvent:
+        if self._event_max_size is not None:
+            event._event_max_size = self._event_max_size
         self.workflow._store.append_event(event)
         self._background_thread_queue.put(event)
         return event
 
     def _run_work_item(self, node: BaseNode[StateType], span_id: UUID) -> None:
+        for event in self.run_node(node, span_id):
+            self._workflow_event_inner_queue.put(event)
+
+    def run_node(
+        self,
+        node: "BaseNode[StateType]",
+        span_id: UUID,
+    ) -> Generator[NodeEvent, None, None]:
+        """
+        Execute a single node and yield workflow events.
+
+        Args:
+            node: The node instance to execute
+            span_id: Unique identifier for this node execution
+
+        Yields:
+            NodeExecutionEvent: Events emitted during node execution (initiated, streaming, fulfilled, rejected)
+        """
         execution = get_execution_context()
-        self._workflow_event_inner_queue.put(
-            NodeExecutionInitiatedEvent(
-                trace_id=execution.trace_id,
-                span_id=span_id,
-                body=NodeExecutionInitiatedBody(
-                    node_definition=node.__class__,
-                    inputs=node._inputs,
-                ),
-                parent=execution.parent_context,
-            )
+
+        node_output_mocks_map = self.workflow.context.node_output_mocks_map
+
+        yield NodeExecutionInitiatedEvent(
+            trace_id=execution.trace_id,
+            span_id=span_id,
+            body=NodeExecutionInitiatedBody(
+                node_definition=node.__class__,
+                inputs=node._inputs,
+            ),
+            parent=execution.parent_context,
         )
 
         logger.debug(f"Started running node: {node.__class__.__name__}")
@@ -208,8 +531,10 @@ class WorkflowRunner(Generic[StateType]):
             )
             node_run_response: NodeRunResponse
             was_mocked: Optional[bool] = None
-            mock_candidates = self.workflow.context.node_output_mocks_map.get(node.Outputs) or []
+            mock_candidates = node_output_mocks_map.get(node.Outputs) or []
             for mock_candidate in mock_candidates:
+                if mock_candidate.disabled:
+                    continue
                 if mock_candidate.when_condition.resolve(node.state):
                     node_run_response = mock_candidate.then_outputs
                     was_mocked = True
@@ -238,8 +563,9 @@ class WorkflowRunner(Generic[StateType]):
                 streaming_output_queues: Dict[str, Queue] = {}
                 outputs = node.Outputs()
 
-                def initiate_node_streaming_output(output: BaseOutput) -> None:
-                    execution = get_execution_context()
+                def initiate_node_streaming_output(
+                    output: BaseOutput,
+                ) -> Generator[NodeExecutionStreamingEvent, None, None]:
                     streaming_output_queues[output.name] = Queue()
                     output_descriptor = OutputReference(
                         name=output.name,
@@ -251,130 +577,249 @@ class WorkflowRunner(Generic[StateType]):
                         node.state.meta.node_outputs[output_descriptor] = streaming_output_queues[output.name]
                     initiated_output: BaseOutput = BaseOutput(name=output.name)
                     initiated_ports = initiated_output > ports
-                    self._workflow_event_inner_queue.put(
-                        NodeExecutionStreamingEvent(
-                            trace_id=execution.trace_id,
-                            span_id=span_id,
-                            body=NodeExecutionStreamingBody(
-                                node_definition=node.__class__,
-                                output=initiated_output,
-                                invoked_ports=initiated_ports,
-                            ),
-                            parent=execution.parent_context,
+                    yield NodeExecutionStreamingEvent(
+                        trace_id=execution.trace_id,
+                        span_id=span_id,
+                        body=NodeExecutionStreamingBody(
+                            node_definition=node.__class__,
+                            output=initiated_output,
+                            invoked_ports=initiated_ports,
                         ),
+                        parent=execution.parent_context,
                     )
 
                 with execution_context(parent_context=updated_parent_context, trace_id=execution.trace_id):
                     for output in node_run_response:
+                        try:
+                            default_serializer(output)
+                        except (TypeError, ValueError) as exc:
+                            raise NodeException(
+                                message=(
+                                    f"Node {node.__class__.__name__} produced output: "
+                                    f"'{output.name}' that could not be serialized to JSON: {exc}"
+                                ),
+                                code=WorkflowErrorCode.INVALID_OUTPUTS,
+                            ) from exc
                         invoked_ports = output > ports
                         if output.is_initiated:
-                            initiate_node_streaming_output(output)
+                            yield from initiate_node_streaming_output(output)
                         elif output.is_streaming:
                             if output.name not in streaming_output_queues:
-                                initiate_node_streaming_output(output)
+                                yield from initiate_node_streaming_output(output)
 
                             streaming_output_queues[output.name].put(output.delta)
-                            self._workflow_event_inner_queue.put(
-                                NodeExecutionStreamingEvent(
-                                    trace_id=execution.trace_id,
-                                    span_id=span_id,
-                                    body=NodeExecutionStreamingBody(
-                                        node_definition=node.__class__,
-                                        output=output,
-                                        invoked_ports=invoked_ports,
-                                    ),
-                                    parent=execution.parent_context,
+                            yield NodeExecutionStreamingEvent(
+                                trace_id=execution.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=output,
+                                    invoked_ports=invoked_ports,
                                 ),
+                                parent=execution.parent_context,
                             )
                         elif output.is_fulfilled:
                             if output.name in streaming_output_queues:
                                 streaming_output_queues[output.name].put(undefined)
 
                             setattr(outputs, output.name, output.value)
-                            self._workflow_event_inner_queue.put(
-                                NodeExecutionStreamingEvent(
-                                    trace_id=execution.trace_id,
-                                    span_id=span_id,
-                                    body=NodeExecutionStreamingBody(
-                                        node_definition=node.__class__,
-                                        output=output,
-                                        invoked_ports=invoked_ports,
-                                    ),
-                                    parent=execution.parent_context,
-                                )
+                            yield NodeExecutionStreamingEvent(
+                                trace_id=execution.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=output,
+                                    invoked_ports=invoked_ports,
+                                ),
+                                parent=execution.parent_context,
                             )
 
-            node.state.meta.node_execution_cache.fulfill_node_execution(node.__class__, span_id)
+            for descriptor, output_value in outputs:
+                if output_value is undefined:
+                    continue
+                try:
+                    default_serializer(output_value)
+                except (TypeError, ValueError) as exc:
+                    raise NodeException(
+                        message=(
+                            f"Node {node.__class__.__name__} produced output '{descriptor.name}' "
+                            f"that could not be serialized to JSON: {exc}"
+                        ),
+                        code=WorkflowErrorCode.INVALID_OUTPUTS,
+                    ) from exc
 
-            with node.state.__atomic__():
-                for descriptor, output_value in outputs:
-                    if output_value is undefined:
-                        if descriptor in node.state.meta.node_outputs:
-                            del node.state.meta.node_outputs[descriptor]
-                        continue
+            node.state.meta.node_execution_cache.fulfill_node_execution(node.__class__, span_id)
 
-                    node.state.meta.node_outputs[descriptor] = output_value
+            with execution_context(parent_context=updated_parent_context, trace_id=execution.trace_id):
+                with node.state.__atomic__():
+                    for descriptor, output_value in outputs:
+                        if output_value is undefined:
+                            if descriptor in node.state.meta.node_outputs:
+                                del node.state.meta.node_outputs[descriptor]
+                            continue
+                        node.state.meta.node_outputs[descriptor] = output_value
 
-            invoked_ports = ports(outputs, node.state)
-            self._workflow_event_inner_queue.put(
-                NodeExecutionFulfilledEvent(
-                    trace_id=execution.trace_id,
-                    span_id=span_id,
-                    body=NodeExecutionFulfilledBody(
-                        node_definition=node.__class__,
-                        outputs=outputs,
-                        invoked_ports=invoked_ports,
-                        mocked=was_mocked,
-                    ),
-                    parent=execution.parent_context,
-                )
+            try:
+                invoked_ports = ports(outputs, node.state)
+            except NodeException as e:
+                raise NodeException(
+                    message=e.message,
+                    code=e.code,
+                    raw_data={**(e.raw_data or {}), "outputs": outputs.__vellum_encode__()},
+                    stacktrace=e.stacktrace,
+                ) from e
+            yield NodeExecutionFulfilledEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionFulfilledBody(
+                    node_definition=node.__class__,
+                    outputs=outputs,
+                    invoked_ports=invoked_ports,
+                    mocked=was_mocked,
+                ),
+                parent=execution.parent_context,
             )
         except NodeException as e:
-            logger.info(e)
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
+            yield self._handle_run_node_exception(e, "Node Exception", execution, span_id, node)
+        except WorkflowInitializationException as e:
+            yield self._handle_run_node_exception(e, "Workflow Initialization Exception", execution, span_id, node)
+        except InvalidExpressionException as e:
+            yield self._handle_run_node_exception(e, "Invalid Expression Exception", execution, span_id, node)
+        except ApiError as e:
+            captured_stacktrace = traceback.format_exc()
+            # Handle structured 403 credential error responses with integration details
+            # The Django API returns {"message": "...", "integration": {...}} for unresolvable credentials
+            # We detect this by shape (403 + integration field present) rather than a code field
+            if e.status_code == 403 and isinstance(e.body, dict) and e.body.get("integration"):
+                error_message = e.body.get(
+                    "message", "You must authenticate with this integration before you can execute this tool."
+                )
+                raw_data = {"integration": e.body["integration"]}
+                yield NodeExecutionRejectedEvent(
                     trace_id=execution.trace_id,
                     span_id=span_id,
                     body=NodeExecutionRejectedBody(
                         node_definition=node.__class__,
-                        error=e.error,
+                        error=WorkflowError(
+                            message=error_message,
+                            code=WorkflowErrorCode.INTEGRATION_CREDENTIALS_UNAVAILABLE,
+                            raw_data=raw_data,
+                        ),
+                        stacktrace=captured_stacktrace,
                     ),
                     parent=execution.parent_context,
                 )
-            )
-        except WorkflowInitializationException as e:
-            logger.info(e)
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
+            # Handle structured 400 error responses with integration details as INVALID_INPUTS
+            elif e.status_code == 400 and isinstance(e.body, dict) and e.body.get("integration"):
+                error_message = e.body.get("message", "Invalid request to integration.")
+                raw_data = {"integration": e.body["integration"]}
+                yield NodeExecutionRejectedEvent(
                     trace_id=execution.trace_id,
                     span_id=span_id,
                     body=NodeExecutionRejectedBody(
                         node_definition=node.__class__,
-                        error=e.error,
+                        error=WorkflowError(
+                            message=error_message,
+                            code=WorkflowErrorCode.INVALID_INPUTS,
+                            raw_data=raw_data,
+                        ),
+                        stacktrace=captured_stacktrace,
                     ),
                     parent=execution.parent_context,
                 )
-            )
-        except Exception as e:
-            logger.exception(f"An unexpected error occurred while running node {node.__class__.__name__}")
+            else:
+                # For all other ApiErrors, use the existing generic exception behavior
+                error_message = self._parse_error_message(e)
+                if error_message is None:
+                    logger.exception(f"An unexpected error occurred while running node {node.__class__.__name__}")
+                    error_code = WorkflowErrorCode.INTERNAL_ERROR
+                    error_message = "Internal error"
+                else:
+                    error_code = WorkflowErrorCode.NODE_EXECUTION
 
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
+                yield NodeExecutionRejectedEvent(
                     trace_id=execution.trace_id,
                     span_id=span_id,
                     body=NodeExecutionRejectedBody(
                         node_definition=node.__class__,
                         error=WorkflowError(
-                            message=str(e),
-                            code=WorkflowErrorCode.INTERNAL_ERROR,
+                            message=error_message,
+                            code=error_code,
                         ),
+                        stacktrace=captured_stacktrace,
                     ),
                     parent=execution.parent_context,
+                )
+        except Exception as e:
+            error_message = self._parse_error_message(e)
+            captured_stacktrace = traceback.format_exc()
+            if error_message is None:
+                logger.exception(f"An unexpected error occurred while running node {node.__class__.__name__}")
+                error_code = WorkflowErrorCode.INTERNAL_ERROR
+                error_message = "Internal error"
+            else:
+                error_code = WorkflowErrorCode.NODE_EXECUTION
+
+            yield NodeExecutionRejectedEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionRejectedBody(
+                    node_definition=node.__class__,
+                    error=WorkflowError(
+                        message=error_message,
+                        code=error_code,
+                    ),
+                    stacktrace=captured_stacktrace,
                 ),
+                parent=execution.parent_context,
             )
 
         logger.debug(f"Finished running node: {node.__class__.__name__}")
 
+    def _handle_run_node_exception(
+        self,
+        exception: Union[NodeException, WorkflowInitializationException, InvalidExpressionException],
+        prefix: str,
+        execution: ExecutionContext,
+        span_id: UUID,
+        node: BaseNode[StateType],
+    ) -> NodeExecutionRejectedEvent:
+        logger.info(f"{prefix}: {exception}")
+        captured_stacktrace = traceback.format_exc()
+
+        return NodeExecutionRejectedEvent(
+            trace_id=execution.trace_id,
+            span_id=span_id,
+            body=NodeExecutionRejectedBody(
+                node_definition=node.__class__,
+                error=exception.error,
+                stacktrace=captured_stacktrace,
+            ),
+            parent=execution.parent_context,
+        )
+
+    def _parse_error_message(self, exception: Exception) -> Optional[str]:
+        try:
+            _, _, tb = sys.exc_info()
+            if tb:
+                tb_list = traceback.extract_tb(tb)
+                if tb_list:
+                    for frame in reversed(tb_list):
+                        exception_module = next(
+                            (
+                                mod.__name__
+                                for mod in sys.modules.values()
+                                if hasattr(mod, "__file__") and mod.__file__ == frame.filename
+                            ),
+                            None,
+                        )
+                        if exception_module and not exception_module.startswith("vellum."):
+                            return str(exception)
+        except Exception:
+            pass
+
+        return None
+
     def _context_run_work_item(
         self,
         node: BaseNode[StateType],
@@ -398,9 +843,7 @@ class WorkflowRunner(Generic[StateType]):
             return
 
         for port in ports:
-            nodes_to_fork = []
             for edge in port.edges:
-                nodes_to_fork.append(edge.to_node)
                 if port.fork_state:
                     next_state = deepcopy(state)
                     self._state_forks.add(next_state)
@@ -411,21 +854,6 @@ class WorkflowRunner(Generic[StateType]):
                     self._concurrency_queue.put((next_state, edge.to_node, invoked_by))
                 else:
                     self._run_node_if_ready(next_state, edge.to_node, invoked_by)
-                if invoked_by:
-                    previous_node = state.meta.node_execution_cache.__node_execution_lookup__.get(invoked_by)
-                    if previous_node:
-                        previous_node_fork_id = state.meta.node_execution_cache.__node_to_fork_id__.get(previous_node)
-                        if previous_node_fork_id:
-                            state.meta.node_execution_cache.__node_to_fork_id__[edge.to_node] = previous_node_fork_id
-                            # Remove previous node fork in __node_to_fork_id__
-                            state.meta.node_execution_cache.__node_to_fork_id__.pop(previous_node, None)
-
-            # If there are multiple edges, we need to create a fork ID
-            if len(nodes_to_fork) > 1:
-                fork_id = uuid4()
-                state.meta.node_execution_cache.__all_fork_ids__.add(fork_id)
-                for node in nodes_to_fork:
-                    state.meta.node_execution_cache.__node_to_fork_id__[node] = fork_id
 
         if self._max_concurrency:
             num_nodes_to_run = self._max_concurrency - len(self._active_nodes_by_execution_id)
@@ -453,8 +881,24 @@ class WorkflowRunner(Generic[StateType]):
 
             all_deps = self._dependencies[node_class]
             node_span_id = node_class.Trigger._queue_node_execution(state, all_deps, invoked_by)
-            if not node_class.Trigger.should_initiate(state, all_deps, node_span_id):
-                return
+
+            try:
+                if not node_class.Trigger.should_initiate(state, all_deps, node_span_id):
+                    return
+            except NodeException as e:
+                execution = get_execution_context()
+
+                self._workflow_event_outer_queue.put(
+                    WorkflowExecutionRejectedEvent(
+                        trace_id=execution.trace_id,
+                        span_id=node_span_id,
+                        body=WorkflowExecutionRejectedBody(
+                            workflow_definition=self.workflow.__class__,
+                            error=e.error,
+                        ),
+                    )
+                )
+                raise e
 
             execution = get_execution_context()
             node = node_class(state=state, context=self.workflow.context)
@@ -472,7 +916,7 @@ class WorkflowRunner(Generic[StateType]):
             )
             worker_thread.start()
 
-    def _handle_work_item_event(self, event: WorkflowEvent) -> Optional[WorkflowError]:
+    def _handle_work_item_event(self, event: WorkflowEvent) -> Optional[NodeExecutionRejectedEvent]:
         active_node = self._active_nodes_by_execution_id.get(event.span_id)
         if not active_node:
             return None
@@ -480,10 +924,23 @@ class WorkflowRunner(Generic[StateType]):
         node = active_node.node
         if event.name == "node.execution.rejected":
             self._active_nodes_by_execution_id.pop(event.span_id)
-            return event.error
+            return event
 
         if event.name == "node.execution.streaming":
             for workflow_output_descriptor in self.workflow.Outputs:
+                if node.__directly_emit_workflow_output__(event, workflow_output_descriptor):
+                    active_node.was_outputs_streamed = True
+                    self._workflow_event_outer_queue.put(
+                        self._stream_workflow_event(
+                            BaseOutput(
+                                name=workflow_output_descriptor.name,
+                                value=event.output.value,
+                                delta=event.output.delta,
+                            )
+                        )
+                    )
+                    continue
+
                 node_output_descriptor = workflow_output_descriptor.instance
                 if not isinstance(node_output_descriptor, OutputReference):
                     continue
@@ -535,16 +992,83 @@ class WorkflowRunner(Generic[StateType]):
 
         return None
 
+    def _emit_node_cancellation_events(
+        self,
+        error_message: str,
+    ) -> None:
+        """
+        Emit node cancellation events for all active nodes.
+
+        Args:
+            error_message: The error message to include in the cancellation events
+        """
+        parent_context = WorkflowParentContext(
+            span_id=self._initial_state.meta.span_id,
+            workflow_definition=self.workflow.__class__,
+            parent=self._execution_context.parent_context,
+        )
+        captured_stacktrace = "".join(traceback.format_stack())
+        active_span_ids = list(self._active_nodes_by_execution_id.keys())
+        for span_id in active_span_ids:
+            active_node = self._active_nodes_by_execution_id.pop(span_id, None)
+            if active_node is not None:
+                active_node.node.__cancel__(error_message)
+
+                rejection_event = NodeExecutionRejectedEvent(
+                    trace_id=self._execution_context.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionRejectedBody(
+                        node_definition=active_node.node.__class__,
+                        error=WorkflowError(
+                            code=WorkflowErrorCode.NODE_CANCELLED,
+                            message=error_message,
+                        ),
+                        stacktrace=captured_stacktrace,
+                    ),
+                    parent=parent_context,
+                )
+                self._workflow_event_outer_queue.put(rejection_event)
+
     def _initiate_workflow_event(self) -> WorkflowExecutionInitiatedEvent:
+        links: Optional[List[SpanLink]] = None
+
+        if self._span_link_info:
+            previous_trace_id, previous_span_id, root_trace_id, root_span_id = self._span_link_info
+            links = [
+                SpanLink(
+                    trace_id=previous_trace_id,
+                    type="PREVIOUS_SPAN",
+                    span_context=WorkflowParentContext(
+                        workflow_definition=self.workflow.__class__,
+                        span_id=previous_span_id,
+                    ),
+                ),
+                SpanLink(
+                    trace_id=root_trace_id,
+                    type="ROOT_SPAN",
+                    span_context=WorkflowParentContext(
+                        workflow_definition=self.workflow.__class__,
+                        span_id=root_span_id,
+                    ),
+                ),
+            ]
+
+        # Get raw inputs from trigger event data if available
+        # This ensures trigger attributes are included in the serialized inputs
+        raw_inputs = self._trigger._event_data if self._trigger else None
+
         return WorkflowExecutionInitiatedEvent(
             trace_id=self._execution_context.trace_id,
             span_id=self._initial_state.meta.span_id,
             body=WorkflowExecutionInitiatedBody(
                 workflow_definition=self.workflow.__class__,
-                inputs=self._initial_state.meta.workflow_inputs,
+                inputs=self._initial_state.meta.workflow_inputs or BaseInputs(),
                 initial_state=deepcopy(self._initial_state) if self._should_emit_initial_state else None,
+                trigger=self._trigger.__class__ if self._trigger else None,
+                raw_inputs=raw_inputs,
             ),
             parent=self._execution_context.parent_context,
+            links=links,
         )
 
     def _stream_workflow_event(self, output: BaseOutput) -> WorkflowExecutionStreamingEvent:
@@ -558,24 +1082,36 @@ class WorkflowRunner(Generic[StateType]):
             parent=self._execution_context.parent_context,
         )
 
-    def _fulfill_workflow_event(self, outputs: OutputsType) -> WorkflowExecutionFulfilledEvent:
+    def _fulfill_workflow_event(self, outputs: OutputsType, final_state: StateType) -> WorkflowExecutionFulfilledEvent:
         return WorkflowExecutionFulfilledEvent(
             trace_id=self._execution_context.trace_id,
             span_id=self._initial_state.meta.span_id,
             body=WorkflowExecutionFulfilledBody(
                 workflow_definition=self.workflow.__class__,
                 outputs=outputs,
+                final_state=final_state,
             ),
             parent=self._execution_context.parent_context,
         )
 
-    def _reject_workflow_event(self, error: WorkflowError) -> WorkflowExecutionRejectedEvent:
+    def _reject_workflow_event(
+        self, error: WorkflowError, captured_stacktrace: Optional[str] = None
+    ) -> WorkflowExecutionRejectedEvent:
+        if captured_stacktrace is None:
+            try:
+                captured_stacktrace = traceback.format_exc()
+                if captured_stacktrace.strip() == "NoneType: None":
+                    captured_stacktrace = None
+            except Exception:
+                pass
+
         return WorkflowExecutionRejectedEvent(
             trace_id=self._execution_context.trace_id,
             span_id=self._initial_state.meta.span_id,
             body=WorkflowExecutionRejectedBody(
                 workflow_definition=self.workflow.__class__,
                 error=error,
+                stacktrace=captured_stacktrace,
             ),
             parent=self._execution_context.parent_context,
         )
@@ -601,10 +1137,19 @@ class WorkflowRunner(Generic[StateType]):
             parent=self._execution_context.parent_context,
         )
 
+    def _wrapped_stream(self) -> None:
+        """Wrapper for _stream that adds httpx logger span ID context."""
+        with self._httpx_logger_with_span_id():
+            self._stream()
+
     def _stream(self) -> None:
         for edge in self.workflow.get_edges():
             self._dependencies[edge.to_node].add(edge.from_port.node_class)
 
+        # Call trigger initiated hook so nodes can reference trigger state
+        if self._trigger is not None:
+            self._trigger.__on_workflow_initiated__(self._initial_state)
+
         current_parent = WorkflowParentContext(
             span_id=self._initial_state.meta.span_id,
             workflow_definition=self.workflow.__class__,
@@ -619,22 +1164,26 @@ class WorkflowRunner(Generic[StateType]):
                 else:
                     self._concurrency_queue.put((self._initial_state, node_cls, None))
             except NodeException as e:
-                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
+                captured_stacktrace = traceback.format_exc()
+                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error, captured_stacktrace))
                 return
             except WorkflowInitializationException as e:
-                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
+                captured_stacktrace = traceback.format_exc()
+                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error, captured_stacktrace))
                 return
             except Exception:
                 err_message = f"An unexpected error occurred while initializing node {node_cls.__name__}"
                 logger.exception(err_message)
+                captured_stacktrace = traceback.format_exc()
                 self._workflow_event_outer_queue.put(
                     self._reject_workflow_event(
                         WorkflowError(code=WorkflowErrorCode.INTERNAL_ERROR, message=err_message),
+                        captured_stacktrace,
                     )
                 )
                 return
 
-        rejection_error: Optional[WorkflowError] = None
+        rejection_event: Optional[NodeExecutionRejectedEvent] = None
 
         while True:
             if not self._active_nodes_by_execution_id:
@@ -645,9 +1194,13 @@ class WorkflowRunner(Generic[StateType]):
             self._workflow_event_outer_queue.put(event)
 
             with execution_context(parent_context=current_parent, trace_id=self._execution_context.trace_id):
-                rejection_error = self._handle_work_item_event(event)
+                rejection_event = self._handle_work_item_event(event)
 
-            if rejection_error:
+            if rejection_event:
+                failed_node_name = rejection_event.body.node_definition.__name__
+                self._emit_node_cancellation_events(
+                    error_message=f"Node execution cancelled due to {failed_node_name} failure",
+                )
                 break
 
         # Handle any remaining events
@@ -656,9 +1209,11 @@ class WorkflowRunner(Generic[StateType]):
                 self._workflow_event_outer_queue.put(event)
 
                 with execution_context(parent_context=current_parent, trace_id=self._execution_context.trace_id):
-                    rejection_error = self._handle_work_item_event(event)
+                    new_rejection_event = self._handle_work_item_event(event)
 
-                if rejection_error:
+                if new_rejection_event:
+                    if rejection_event is None:
+                        rejection_event = new_rejection_event
                     break
         except Empty:
             pass
@@ -678,10 +1233,24 @@ class WorkflowRunner(Generic[StateType]):
             )
             return
 
-        if rejection_error:
-            self._workflow_event_outer_queue.put(self._reject_workflow_event(rejection_error))
+        if rejection_event:
+            self._workflow_event_outer_queue.put(
+                self._reject_workflow_event(rejection_event.error, rejection_event.body.stacktrace)
+            )
             return
 
+        # Call trigger fulfilled hook before output resolution so state changes are reflected in outputs
+        if self._trigger is not None:
+            self._trigger.__on_workflow_fulfilled__(final_state)
+
+            # Drain any events produced by the trigger hook (e.g., snapshot events from state mutations)
+            # and forward them to the outer queue so stream consumers can observe them
+            try:
+                while event := self._workflow_event_inner_queue.get_nowait():
+                    self._workflow_event_outer_queue.put(event)
+            except Empty:
+                pass
+
         fulfilled_outputs = self.workflow.Outputs()
         for descriptor, value in fulfilled_outputs:
             if isinstance(value, BaseDescriptor):
@@ -693,7 +1262,7 @@ class WorkflowRunner(Generic[StateType]):
                     descriptor.instance.resolve(final_state),
                 )
 
-        self._workflow_event_outer_queue.put(self._fulfill_workflow_event(fulfilled_outputs))
+        self._workflow_event_outer_queue.put(self._fulfill_workflow_event(fulfilled_outputs, final_state))
 
     def _run_background_thread(self) -> None:
         state_class = self.workflow.get_state_class()
@@ -715,16 +1284,44 @@ class WorkflowRunner(Generic[StateType]):
 
         while not kill_switch.wait(timeout=0.1):
             if self._cancel_signal.is_set():
+                self._emit_node_cancellation_events(
+                    error_message="Workflow run cancelled",
+                )
+
+                captured_stacktrace = "".join(traceback.format_stack())
                 self._workflow_event_outer_queue.put(
                     self._reject_workflow_event(
                         WorkflowError(
                             code=WorkflowErrorCode.WORKFLOW_CANCELLED,
                             message="Workflow run cancelled",
-                        )
+                        ),
+                        captured_stacktrace,
                     )
                 )
                 return
 
+    def _run_timeout_thread(self, kill_switch: ThreadingEvent) -> None:
+        if not self._timeout:
+            return
+
+        if kill_switch.wait(timeout=self._timeout):
+            return
+
+        self._emit_node_cancellation_events(
+            error_message=f"Workflow execution exceeded timeout of {self._timeout} seconds",
+        )
+
+        captured_stacktrace = "".join(traceback.format_stack())
+        self._workflow_event_outer_queue.put(
+            self._reject_workflow_event(
+                WorkflowError(
+                    code=WorkflowErrorCode.WORKFLOW_TIMEOUT,
+                    message=f"Workflow execution exceeded timeout of {self._timeout} seconds",
+                ),
+                captured_stacktrace,
+            )
+        )
+
     def _is_terminal_event(self, event: WorkflowEvent) -> bool:
         if (
             event.name == "workflow.execution.fulfilled"
@@ -734,21 +1331,30 @@ class WorkflowRunner(Generic[StateType]):
             return event.workflow_definition == self.workflow.__class__
         return False
 
-    def stream(self) -> WorkflowEventStream:
-        background_thread = Thread(
+    def _generate_events(self) -> Generator[WorkflowEvent, None, None]:
+        self._background_thread = Thread(
             target=self._run_background_thread,
             name=f"{self.workflow.__class__.__name__}.background_thread",
         )
-        background_thread.start()
+        self._background_thread.start()
 
         cancel_thread_kill_switch = ThreadingEvent()
         if self._cancel_signal:
-            cancel_thread = Thread(
+            self._cancel_thread = Thread(
                 target=self._run_cancel_thread,
                 name=f"{self.workflow.__class__.__name__}.cancel_thread",
                 kwargs={"kill_switch": cancel_thread_kill_switch},
             )
-            cancel_thread.start()
+            self._cancel_thread.start()
+
+        timeout_thread_kill_switch = ThreadingEvent()
+        if self._timeout:
+            self._timeout_thread = Thread(
+                target=self._run_timeout_thread,
+                name=f"{self.workflow.__class__.__name__}.timeout_thread",
+                kwargs={"kill_switch": timeout_thread_kill_switch},
+            )
+            self._timeout_thread.start()
 
         event: WorkflowEvent
         if self._is_resuming:
@@ -759,13 +1365,13 @@ class WorkflowRunner(Generic[StateType]):
         yield self._emit_event(event)
 
         # The extra level of indirection prevents the runner from waiting on the caller to consume the event stream
-        stream_thread = Thread(
-            target=self._stream,
+        self._stream_thread = Thread(
+            target=self._wrapped_stream,
             name=f"{self.workflow.__class__.__name__}.stream_thread",
         )
-        stream_thread.start()
+        self._stream_thread.start()
 
-        while stream_thread.is_alive():
+        while self._stream_thread.is_alive():
             try:
                 event = self._workflow_event_outer_queue.get(timeout=0.1)
             except Empty:
@@ -792,3 +1398,24 @@ class WorkflowRunner(Generic[StateType]):
 
         self._background_thread_queue.put(None)
         cancel_thread_kill_switch.set()
+        timeout_thread_kill_switch.set()
+
+    def stream(self) -> WorkflowEventStream:
+        return WorkflowEventGenerator(self._generate_events(), self._initial_state.meta.span_id)
+
+    def join(self) -> None:
+        """
+        Wait for all background threads to complete.
+        This ensures all pending work is finished before the runner terminates.
+        """
+        if self._stream_thread and self._stream_thread.is_alive():
+            self._stream_thread.join()
+
+        if self._background_thread and self._background_thread.is_alive():
+            self._background_thread.join()
+
+        if self._cancel_thread and self._cancel_thread.is_alive():
+            self._cancel_thread.join()
+
+        if self._timeout_thread and self._timeout_thread.is_alive():
+            self._timeout_thread.join()