fix span handling and span hierarchy propagation (#381)

.changeset/afraid-terms-kick.md

@@ -0,0 +1,5 @@
+---
+"llama-index-workflows": patch
+---
+
+Fix span tracking in observability tooling

packages/llama-index-workflows/src/workflows/plugins/basic.py

@@ -15,6 +15,7 @@ if TYPE_CHECKING:
     from workflows.workflow import Workflow
 
 from llama_index_instrumentation.dispatcher import active_instrument_tags
+from llama_index_instrumentation.span import active_span_id
 
 from workflows.context.serializers import BaseSerializer, JsonSerializer
 from workflows.context.state_store import (
@@ -284,13 +285,20 @@ class BasicRuntime(Runtime):
         queues = self._get_or_create_queues(run_id, init_state)
         queues.state_store = state_store
 
+        # Capture parent span ID and instrument tags BEFORE creating the task
+        # (they won't be inherited by the background task)
+        captured_tags = {**active_instrument_tags.get()}
+        parent_span_id = active_span_id.get()
+        if parent_span_id is not None:
+            captured_tags["parent_span_id"] = parent_span_id
+
         async def run_with_concurrency_limit() -> StopEvent:
             # Capture strong reference to queues for the task's lifetime,
             # enabling fire-and-forget even if the caller drops the external adapter.
             _ = queues
             async with self._maybe_acquire_max_concurrent_runs(workflow, run_id):
                 return await registered.workflow_run_fn(
-                    init_state, start_event, active_instrument_tags.get()
+                    init_state, start_event, captured_tags
                 )
 
         with setting_run_id(run_id):

packages/llama-index-workflows/src/workflows/runtime/types/plugin.py

@@ -602,5 +602,5 @@ class WorkflowRunFunction(Protocol):
         self,
         init_state: BrokerState,
         start_event: StartEvent | None = None,
-        tags: dict[str, Any] = {},
+        tags: dict[str, Any] | None = None,
     ) -> Coroutine[None, None, StopEvent]: ...

packages/llama-index-workflows/src/workflows/runtime/types/step_function.py

@@ -10,6 +10,7 @@ from contextvars import copy_context
 from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Protocol
 
 from llama_index_instrumentation.dispatcher import instrument_tags
+from llama_index_instrumentation.span import active_span_id
 from workflows.decorators import P, R, StepConfig
 from workflows.errors import WorkflowRuntimeError
 from workflows.events import (
@@ -181,7 +182,7 @@ def create_workflow_run_function(
     async def run_workflow(
         init_state: BrokerState,
         start_event: StartEvent | None = None,
-        tags: dict[str, Any] = {},
+        tags: dict[str, Any] | None = None,
     ) -> StopEvent:
         from workflows.context.context import Context
         from workflows.context.internal_context import InternalContext
@@ -190,28 +191,43 @@ def create_workflow_run_function(
         # Set run_id context before creating internal context
         internal_ctx = Context._create_internal(workflow=workflow)
         internal_adapter = workflow._runtime.get_internal_adapter(workflow)
-        with instrument_tags(tags):
-            # defer execution to make sure the task can be captured and passed
-            # to the handler as async exception, protecting against exceptions from before_start
-            await asyncio.sleep(0)
 
-            run_ctx = RunContext(
-                workflow=workflow,
-                run_adapter=internal_adapter,
-                context=internal_ctx,
-                steps=registered.steps,
-            )
-            try:
-                with run_context(run_ctx):
-                    result = await control_loop_fn(
-                        start_event,
-                        init_state,
-                        internal_adapter.run_id,
-                    )
-                    return result
-            finally:
-                # Cancel any background tasks from InternalContext on completion or cancellation
-                if isinstance(internal_ctx._face, InternalContext):
-                    internal_ctx._face.cancel_background_tasks()
+        # Extract parent span ID if present and remove from tags
+        tags = tags or {}
+        parent_span_id = tags.pop("parent_span_id", None)
+
+        # Set parent span ID context if provided
+        parent_span_token = None
+        if parent_span_id is not None:
+            parent_span_token = active_span_id.set(parent_span_id)
+
+        try:
+            with instrument_tags(tags):
+                # defer execution to make sure the task can be captured and passed
+                # to the handler as async exception, protecting against exceptions from before_start
+                await asyncio.sleep(0)
+
+                run_ctx = RunContext(
+                    workflow=workflow,
+                    run_adapter=internal_adapter,
+                    context=internal_ctx,
+                    steps=registered.steps,
+                )
+                try:
+                    with run_context(run_ctx):
+                        result = await control_loop_fn(
+                            start_event,
+                            init_state,
+                            internal_adapter.run_id,
+                        )
+                        return result
+                finally:
+                    # Cancel any background tasks from InternalContext on completion or cancellation
+                    if isinstance(internal_ctx._face, InternalContext):
+                        internal_ctx._face.cancel_background_tasks()
+        finally:
+            # Reset parent span ID if it was set
+            if parent_span_token is not None:
+                active_span_id.reset(parent_span_token)
 
     return run_workflow

packages/llama-index-workflows/src/workflows/workflow.py

@@ -4,7 +4,9 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import logging
+import uuid
 from dataclasses import dataclass, field
 from typing import (
     TYPE_CHECKING,
@@ -13,6 +15,9 @@ from typing import (
 )
 
 from llama_index_instrumentation import get_dispatcher
+from llama_index_instrumentation.dispatcher import active_instrument_tags
+from llama_index_instrumentation.events.span import SpanDropEvent
+from llama_index_instrumentation.span import active_span_id
 from pydantic import ValidationError
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -355,7 +360,6 @@ class Workflow(metaclass=WorkflowMeta):
             logger.debug(e)
             raise WorkflowRuntimeError(msg)
 
-    @dispatcher.span
     def run(
         self,
         ctx: Context | None = None,
@@ -410,25 +414,92 @@ class Workflow(metaclass=WorkflowMeta):
         """
         from workflows.context import Context
 
-        # Validate the workflow
-        self._validate()
+        # Manually manage span to keep it open until workflow completes
+        # llama-index-instrumentation currently does not manage Awaitable's well (i.e. the workflow handler)
+        # this pattern is unusual enough to special case it here
+        # First, generate span ID
+        cls_name = self.__class__.__name__
+        span_id = f"{cls_name}.run-{uuid.uuid4()}"
 
-        # Extract run_id before passing remaining kwargs to start event
-        run_id = kwargs.pop("run_id", None)
+        # Get parent span ID for nesting
+        parent_span_id = active_span_id.get()
 
-        # If a previous context is provided, pass its serialized form
-        ctx = ctx if ctx is not None else Context(self)
-        # TODO(v3) - remove dependency on is running for choosing whether to send a StartEvent.
-        # Is not an easily synchronously queryable property.
-        start_event_instance: StartEvent | None = (
-            None
-            if ctx.is_running
-            else self._get_start_event_instance(start_event, **kwargs)
+        # Create bound args for span_enter/exit
+        bound_args = inspect.signature(self.run).bind(
+            ctx=ctx, start_event=start_event, **kwargs
         )
-        return ctx._workflow_run(
-            workflow=self, start_event=start_event_instance, run_id=run_id
+
+        # Set active span and notify span handlers
+        span_token = active_span_id.set(span_id)
+
+        dispatcher.span_enter(
+            id_=span_id,
+            bound_args=bound_args,
+            instance=self,
+            parent_id=parent_span_id,
+            tags=active_instrument_tags.get(),
         )
 
+        try:
+            # Validate the workflow
+            self._validate()
+
+            # Extract run_id before passing remaining kwargs to start event
+            run_id = kwargs.pop("run_id", None)
+
+            # If a previous context is provided, pass its serialized form
+            ctx = ctx if ctx is not None else Context(self)
+            # TODO(v3) - remove dependency on is running for choosing whether to send a StartEvent.
+            # Is not an easily synchronously queryable property.
+            start_event_instance: StartEvent | None = (
+                None
+                if ctx.is_running
+                else self._get_start_event_instance(start_event, **kwargs)
+            )
+            handler = ctx._workflow_run(
+                workflow=self, start_event=start_event_instance, run_id=run_id
+            )
+
+            # Add callback to close span when workflow completes
+            def _on_workflow_complete(task: asyncio.Task[Any]) -> None:
+                try:
+                    # Get result or exception
+                    if task.cancelled():
+                        result = None
+                    else:
+                        try:
+                            result = task.result()
+                        except Exception:
+                            result = None
+
+                    # Notify span exit
+                    dispatcher.span_exit(
+                        id_=span_id,
+                        bound_args=bound_args,
+                        instance=self,
+                        result=result,
+                    )
+                finally:
+                    # Reset span token
+                    try:
+                        active_span_id.reset(span_token)
+                    except ValueError:
+                        # Token might be from different context, ignore
+                        pass
+
+            # Attach callback to the handler's result task
+            handler._result_task.add_done_callback(_on_workflow_complete)
+
+            return handler
+        except BaseException as e:
+            # If run() fails, drop the span
+            dispatcher.event(SpanDropEvent(span_id=span_id, err_str=str(e)))
+            dispatcher.span_drop(
+                id_=span_id, bound_args=bound_args, instance=self, err=e
+            )
+            active_span_id.reset(span_token)
+            raise
+
     def _validate_resource_configs(self) -> list[str]:
         """Validate all resource configs (including nested ones) by loading them."""
         errors: list[str] = []