properly start spans

Author: loukratz-bv

qtype/interpreter/base/base_step_executor.py

@@ -97,21 +97,6 @@ def _resolve_secret(self, value: str | SecretReference) -> str:
         context = f"step '{self.step.id}'"
         return self._secret_manager(value, context)
 
-    def _enrich_with_span_metadata(
-        self, message: FlowMessage, span: trace.Span
-    ) -> FlowMessage:
-        """Add span_id and trace_id to message metadata for feedback tracking."""
-        if not span.is_recording():
-            return message
-
-        span_context = span.get_span_context()
-        updated_metadata = {
-            **message.metadata,
-            "span_id": format(span_context.span_id, "016x"),
-            "trace_id": format(span_context.trace_id, "032x"),
-        }
-        return message.model_copy(update={"metadata": updated_metadata})
-
     async def _filter_and_collect_errors(
         self,
         message_stream: AsyncIterator[FlowMessage],
@@ -191,7 +176,10 @@ async def execute(
             Processed messages, with failed messages emitted first
         """
         # Start a span for tracking
-        # Note: We manually manage the span lifecycle to allow yielding
+        # Note: We do NOT attach this span to context here to avoid
+        # making upstream steps children of this step when we consume
+        # the input stream. Instead, _process_message_with_telemetry
+        # will attach it when calling process_message().
         span = self._tracer.start_span(
             f"step.{self.step.id}",
             attributes={
@@ -201,10 +189,8 @@ async def execute(
             },
         )
 
-        # Make this span the active context so child spans will nest under it
-        # Only attach if span is recording (i.e., real tracer is configured)
-        ctx = trace.set_span_in_context(span)
-        token = context.attach(ctx) if span.is_recording() else None
+        # Store span in self so _process_message_with_telemetry can access it
+        self._current_step_span = span
 
         # Initialize the cache
         # this is done once per execution so re-runs are fast
@@ -259,7 +245,7 @@ async def process_item(
                         self.progress.update_for_message(
                             result, self.context.on_progress
                         )
-                        yield self._enrich_with_span_metadata(result, span)
+                        yield result
 
                 # Emit failed messages after processing completes
                 for msg in failed_messages:
@@ -268,14 +254,14 @@ async def process_item(
                     self.progress.update_for_message(
                         msg, self.context.on_progress
                     )
-                    yield self._enrich_with_span_metadata(msg, span)
+                    yield msg
 
                 # Finalize and track those messages too
                 async for msg in self.finalize():
                     message_count += 1
                     if msg.is_failed():
                         error_count += 1
-                    yield self._enrich_with_span_metadata(msg, span)
+                    yield msg
 
                 # Close the cache
                 if self.cache:
@@ -302,10 +288,8 @@ async def process_item(
                 span.set_status(Status(StatusCode.ERROR, f"Step failed: {e}"))
                 raise
             finally:
-                # Detach the context and end the span
-                # Only detach if we successfully attached (span was recording)
-                if token is not None:
-                    context.detach(token)
+                # Clean up step span reference and end the span
+                self._current_step_span = None
                 span.end()
 
     @abstractmethod
@@ -383,51 +367,77 @@ async def _process_message_with_telemetry(
 
         This method creates a child span for each message processing
         operation, automatically recording errors and success metrics.
-        The child span will automatically be nested under the current
-        active span in the context.
+        The child span will be nested under the step span.
+
+        The step span context is attached here (not in execute()) to
+        ensure step spans are siblings under the flow span, not nested.
         """
-        # Get current context and create child span within it
-        span = self._tracer.start_span(
-            f"step.{self.step.id}.process_message",
-            attributes={
-                "session.id": message.session.session_id,
-            },
-        )
+        # Attach step span context so process_message span becomes its child
+        step_span = getattr(self, "_current_step_span", None)
+        if step_span and step_span.is_recording():
+            ctx = trace.set_span_in_context(step_span)
+            token = context.attach(ctx)
+        else:
+            token = None
 
         try:
-            output_count = 0
-            error_occurred = False
-
-            async for output_msg in self.process_message(message):
-                output_count += 1
-                if output_msg.is_failed():
-                    error_occurred = True
-                    span.add_event(
-                        "message_failed",
-                        {
-                            "error": str(output_msg.error),
-                        },
+            # Create child span for this specific message processing
+            span = self._tracer.start_span(
+                f"step.{self.step.id}.process_message",
+                attributes={
+                    "session.id": message.session.session_id,
+                },
+            )
+
+            try:
+                output_count = 0
+                error_occurred = False
+
+                async for output_msg in self.process_message(message):
+                    output_count += 1
+                    if output_msg.is_failed():
+                        error_occurred = True
+                        span.add_event(
+                            "message_failed",
+                            {
+                                "error": str(output_msg.error),
+                            },
+                        )
+                    # Enrich with process_message span for feedback tracking
+                    span_context = span.get_span_context()
+                    updated_metadata = {
+                        **output_msg.metadata,
+                        "span_id": format(span_context.span_id, "016x"),
+                        "trace_id": format(span_context.trace_id, "032x"),
+                    }
+                    yield output_msg.model_copy(
+                        update={"metadata": updated_metadata}
                     )
-                yield output_msg
 
-            # Record processing metrics
-            span.set_attribute("message.outputs", output_count)
+                # Record processing metrics
+                span.set_attribute("message.outputs", output_count)
 
-            if error_occurred:
+                if error_occurred:
+                    span.set_status(
+                        Status(
+                            StatusCode.ERROR, "Message processing had errors"
+                        )
+                    )
+                else:
+                    span.set_status(Status(StatusCode.OK))
+
+            except Exception as e:
+                span.record_exception(e)
                 span.set_status(
-                    Status(StatusCode.ERROR, "Message processing had errors")
+                    Status(StatusCode.ERROR, f"Processing failed: {e}")
                 )
-            else:
-                span.set_status(Status(StatusCode.OK))
-
-        except Exception as e:
-            span.record_exception(e)
-            span.set_status(
-                Status(StatusCode.ERROR, f"Processing failed: {e}")
-            )
-            raise
+                raise
+            finally:
+                span.end()
         finally:
-            span.end()
+            # Detach step span context
+            if token is not None:
+                context.detach(token)
 
     async def finalize(self) -> AsyncIterator[FlowMessage]:
         """

tests/interpreter/test_step_executor.py

@@ -427,9 +427,10 @@ async def test_span_metadata_enrichment(self, simple_step, session):
             assert len(trace_id) == 32
             assert all(c in "0123456789abcdef" for c in trace_id)
 
-        # Span ID should be the same for all messages in the same step
+        # Each message should have a unique process_message span_id
+        # (this allows per-message feedback instead of per-step)
         span_ids = [r.metadata["span_id"] for r in results]
-        assert len(set(span_ids)) == 1
+        assert len(set(span_ids)) == len(results)
 
         # Trace ID should be the same for all messages in the same execution
         trace_ids = [r.metadata["trace_id"] for r in results]