fix(telemetry): correct streaming validation span nesting and cross-task detach

The backend `chat` span opened by stream_with_chunking attaches in the caller
task but finishes in the orchestration task that drains the MOT. Detaching its
OTel context token from that second task failed (logged by OTel as "Failed to
detach context"), and the failed detach left the generation span ambient, so a
subsequent validation `chat` span nested under generation instead of being its
sibling under stream_with_chunking.

Add STREAMING_ORCHESTRATION_START/END hooks fired on the orchestration task.
The tracing plugin uses them to re-attach the stream_with_chunking span as that
task's ambient context for the drain/validate loop, so mid-stream spans parent
correctly. _safe_detach skips a detach that would cross asyncio tasks: within a
reattach scope the skip is expected and logged at debug; otherwise the detach
runs so OTel surfaces its own error, preceded by a warning naming the mismatch.

Add unit coverage for the reattach helpers, the task-identity classification,
and both log paths; an integration test asserting both chat spans are siblings
under stream_with_chunking and carry no streaming events; and an e2e streaming
test. Document the stream_with_chunking span and its hierarchy.

Assisted-by: Claude Code
Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>

Author: ajbozarth

docs/docs/observability/tracing.md

@@ -137,6 +137,23 @@ The `action` span:
 | `mellea.response` | Model response truncated to 500 characters; recorded only when `MELLEA_TRACES_CONTENT=true` |
 | `mellea.response_length` | Length of the model response (always recorded) |
 
+The `stream_with_chunking` span covers one streaming-generation run. It wraps
+the backend generation and any per-chunk validation:
+
+| Attribute | Description |
+| --------- | ----------- |
+| `mellea.has_requirements` | Whether requirements were supplied |
+| `mellea.requirement_count` | Number of requirements supplied |
+| `mellea.chunking_strategy` | `ChunkingStrategy` class name (e.g., `SentenceChunker`) |
+| `mellea.full_text_length` | Length of the accumulated text at completion |
+| `gen_ai.request.model` | Model ID, when known |
+| `gen_ai.provider.name` | Provider name, when known |
+
+It also records span events through the run: `quick_check` and `chunk` per
+validated chunk, `streaming_done` once the stream drains, `full_validation`
+after the final `validate()` calls, `error` on an unhandled exception, and
+`completed` when the run exits.
+
 ### Backend spans (`mellea.backend`)
 
 Backend spans cover individual LLM API calls. They follow the
@@ -189,6 +206,22 @@ session                   (mellea.application)
                           [gen_ai.request.model=gpt-4o]
 ```
 
+In a `stream_with_chunking` run, the backend generation and each per-chunk
+validation call nest under the `stream_with_chunking` span as sibling `chat`
+spans:
+
+```text
+stream_with_chunking      (mellea.application)
+│                         [mellea.chunking_strategy=SentenceChunker]
+├── chat                  (mellea.backend)    ← streaming generation
+│                         [gen_ai.request.model=granite4.1:3b]
+└── chat                  (mellea.backend)    ← per-chunk validation
+                          [gen_ai.request.model=granite4.1:3b]
+```
+
+The `stream_with_chunking` span itself parents under whatever span is active
+when the run starts, or is a root span when none is.
+
 ## Reading traces in a typical agent run
 
 When you open a trace in your backend, look for these patterns:

mellea/plugins/hooks/streaming.py

@@ -67,3 +67,26 @@ class StreamingEndPayload(MelleaBasePayload):
     model: str | None = None
     provider: str | None = None
     full_text_length: int = 0
+
+
+class StreamingOrchestrationStartPayload(MelleaBasePayload):
+    """Payload for `streaming_orchestration_start` — on the orchestration task, before the stream is drained.
+
+    Attributes:
+        streaming_id: UUID correlating with the matching `streaming_start`.
+    """
+
+    streaming_id: str = ""
+
+
+class StreamingOrchestrationEndPayload(MelleaBasePayload):
+    """Payload for `streaming_orchestration_end` — on the orchestration task, after the stream is drained.
+
+    Fires on the same task as `streaming_orchestration_start`.
+
+    Attributes:
+        streaming_id: UUID correlating with the matching
+            `streaming_orchestration_start`.
+    """
+
+    streaming_id: str = ""

mellea/plugins/types.py

@@ -71,6 +71,8 @@ class HookType(StrEnum):
     STREAMING_START = "streaming_start"
     STREAMING_EVENT = "streaming_event"
     STREAMING_END = "streaming_end"
+    STREAMING_ORCHESTRATION_START = "streaming_orchestration_start"
+    STREAMING_ORCHESTRATION_END = "streaming_orchestration_end"
 
 
 # Lazily populated mapping: hook_type -> (payload_class, result_class).
@@ -111,6 +113,8 @@ def _build_hook_registry() -> dict[str, tuple[type, type]]:
     from mellea.plugins.hooks.streaming import (
         StreamingEndPayload,
         StreamingEventPayload,
+        StreamingOrchestrationEndPayload,
+        StreamingOrchestrationStartPayload,
         StreamingStartPayload,
     )
     from mellea.plugins.hooks.tool import ToolPostInvokePayload, ToolPreInvokePayload
@@ -169,6 +173,14 @@ def _build_hook_registry() -> dict[str, tuple[type, type]]:
         HookType.STREAMING_START.value: (StreamingStartPayload, PluginResult),
         HookType.STREAMING_EVENT.value: (StreamingEventPayload, PluginResult),
         HookType.STREAMING_END.value: (StreamingEndPayload, PluginResult),
+        HookType.STREAMING_ORCHESTRATION_START.value: (
+            StreamingOrchestrationStartPayload,
+            PluginResult,
+        ),
+        HookType.STREAMING_ORCHESTRATION_END.value: (
+            StreamingOrchestrationEndPayload,
+            PluginResult,
+        ),
     }
 
 

mellea/stdlib/streaming.py

@@ -551,6 +551,15 @@ async def _process_chunk(c: str, ci: int) -> bool:
         return False
 
     try:
+        # Inside the try so a cancellation at this await still runs the finally.
+        if has_plugins(HookType.STREAMING_ORCHESTRATION_START):
+            from ..plugins.hooks.streaming import StreamingOrchestrationStartPayload
+
+            await invoke_hook(
+                HookType.STREAMING_ORCHESTRATION_START,
+                StreamingOrchestrationStartPayload(streaming_id=result._streaming_id),
+            )
+
         while not mot.is_computed():
             try:
                 delta = await mot.astream()
@@ -680,6 +689,19 @@ async def _process_chunk(c: str, ci: int) -> bool:
             except BaseException:
                 pass
 
+        # Shielded so a CancelledError from the hook cannot skip the terminal
+        # queue bookkeeping and _done.set() below.
+        if has_plugins(HookType.STREAMING_ORCHESTRATION_END):
+            from ..plugins.hooks.streaming import StreamingOrchestrationEndPayload
+
+            try:
+                await invoke_hook(
+                    HookType.STREAMING_ORCHESTRATION_END,
+                    StreamingOrchestrationEndPayload(streaming_id=result._streaming_id),
+                )
+            except BaseException:
+                pass
+
         completed_ev = CompletedEvent(
             success=result.completed, full_text=result.full_text, attempts_used=1
         )

mellea/telemetry/tracing.py

@@ -24,6 +24,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import os
 import warnings
 from importlib.metadata import version
@@ -221,7 +222,75 @@ def get_backend_tracer() -> Any:
     return _backend_tracer
 
 
-_in_flight_spans: dict[str, tuple[Span, Token[Context]]] = {}
+_in_flight_spans: dict[str, tuple[Span, Token[Context], asyncio.Task[Any] | None]] = {}
+
+# reattach_span() entries, keyed by correlation key: the OTel context token plus
+# the task that attached it. Released by release_reattached_span() on that task.
+_reattached_tokens: dict[str, tuple[Token[Context], asyncio.Task[Any] | None]] = {}
+
+
+def _current_task() -> asyncio.Task[Any] | None:
+    """Return the running asyncio task, or None when no loop is running."""
+    try:
+        return asyncio.current_task()
+    except RuntimeError:
+        return None
+
+
+def _safe_detach(token: Token[Context], attach_task: asyncio.Task[Any] | None) -> None:
+    """Detach `token`, suppressing only the cross-task detach we expect and understand.
+
+    OTel context tokens are bound to the `contextvars.Context` of the task that
+    created them, so detaching from a different task fails — OTel catches the
+    `ValueError` and logs it at ERROR as "Failed to detach context". Most spans
+    attach and finish on one task and never hit this.
+
+    A cross-task detach is suppressed (skipped, since it would only fail) and
+    logged at debug only when the detaching task holds an open `reattach_span`
+    scope — the marker that this task knowingly opened a span elsewhere and
+    expects the mismatch. Any other cross-task detach is left to run so OTel
+    surfaces its ERROR with a traceback to the real origin; a warning is added
+    first to name the task mismatch OTel's message omits.
+
+    Example:
+        Under `stream_with_chunking` the backend `chat` span attaches in the
+        caller task but finishes in the orchestration task that drains the MOT.
+        To keep that span's `chat` children nesting correctly, the orchestration
+        task re-attaches the streaming span for the duration of the drain
+        (`reattach_span` / `release_reattached_span`). The cross-task `chat`
+        detach that then happens within that scope is the expected, suppressed
+        case.
+
+    Note:
+        The reattach scope is an *incomplete* proxy for "expected". It holds for
+        streaming because that case both needs sibling-nesting protection (so it
+        reattaches) and has an expected cross-task detach. A future case with the
+        same open-in-parent / close-in-child shape but no siblings to protect
+        would not reattach, so its equally-expected cross-task detach falls
+        through to the warn-and-detach path. That is harmless (the detach only
+        fails, and the task ends right after, so nothing leaks); the warning is
+        the signal that the new use needs its own way to mark the detach expected.
+
+    Args:
+        token: The OTel context token returned by the matching `attach`.
+        attach_task: The task that performed the `attach`, or None if it was
+            attached outside any running task.
+    """
+    current = _current_task()
+    if attach_task is not None and current is not attach_task:
+        from mellea.core.utils import MelleaLogger
+
+        if any(task is current for _, task in _reattached_tokens.values()):
+            MelleaLogger.get_logger().debug(
+                "Skipped expected cross-task OTel context detach within a "
+                "reattached-span scope."
+            )
+            return
+        MelleaLogger.get_logger().warning(
+            "Detaching an OTel context token across asyncio tasks; the span's "
+            "attach and detach ran on different tasks. OTel will log the failure."
+        )
+    otel_context.detach(token)
 
 
 def start_backend_span(
@@ -286,7 +355,7 @@ def start_backend_span(
     set_conversation_id(span)
 
     token = otel_context.attach(trace.set_span_in_context(span))
-    _in_flight_spans[generation_id] = (span, token)
+    _in_flight_spans[generation_id] = (span, token, _current_task())
     return span
 
 
@@ -321,7 +390,7 @@ def finish_backend_span_success(
     entry = _in_flight_spans.pop(generation_id, None)
     if entry is None:
         return
-    span, token = entry
+    span, token, attach_task = entry
     try:
         if gen is not None:
             set_request_attrs(span, gen, operation)
@@ -330,7 +399,7 @@ def finish_backend_span_success(
         if mot is not None and gen is not None:
             set_mellea_attrs(span, mot, gen)
     finally:
-        otel_context.detach(token)
+        _safe_detach(token, attach_task)
         span.end()
 
 
@@ -355,15 +424,15 @@ def finish_backend_span_error(
     entry = _in_flight_spans.pop(generation_id, None)
     if entry is None:
         return
-    span, token = entry
+    span, token, attach_task = entry
     try:
         if gen is not None:
             set_request_attrs(span, gen, operation)
         span.record_exception(exception)
         span.set_status(trace.Status(trace.StatusCode.ERROR, str(exception)))
         span.set_attribute("error.type", type(exception).__name__)
     finally:
-        otel_context.detach(token)
+        _safe_detach(token, attach_task)
         span.end()
 
 
@@ -392,7 +461,7 @@ def _start_application_span(
             set_attribute_safe(span, k, v)
 
     token = otel_context.attach(trace.set_span_in_context(span))
-    _in_flight_spans[key] = (span, token)
+    _in_flight_spans[key] = (span, token, _current_task())
     return span
 
 
@@ -414,13 +483,13 @@ def _finish_application_span_success(
     entry = _in_flight_spans.pop(key, None)
     if entry is None:
         return
-    span, token = entry
+    span, token, attach_task = entry
     try:
         if extra_attributes:
             for k, v in extra_attributes.items():
                 set_attribute_safe(span, k, v)
     finally:
-        otel_context.detach(token)
+        _safe_detach(token, attach_task)
         span.end()
 
 
@@ -448,7 +517,7 @@ def _finish_application_span_error(
     entry = _in_flight_spans.pop(key, None)
     if entry is None:
         return
-    span, token = entry
+    span, token, attach_task = entry
     try:
         if extra_attributes:
             for k, v in extra_attributes.items():
@@ -460,7 +529,7 @@ def _finish_application_span_error(
         else:
             span.set_status(trace.Status(trace.StatusCode.ERROR, description or ""))
     finally:
-        otel_context.detach(token)
+        _safe_detach(token, attach_task)
         span.end()
 
 
@@ -697,11 +766,45 @@ def add_streaming_event(
     entry = _in_flight_spans.get(streaming_id)
     if entry is None:
         return
-    span, _ = entry
+    span = entry[0]
     filtered = {k: v for k, v in attributes.items() if v is not None}
     span.add_event(event_name, filtered)
 
 
+def reattach_span(key: str) -> None:
+    """Make the in-flight span `key` the current task's ambient context.
+
+    Spans opened by later work on this task then parent under it. Paired with
+    `release_reattached_span()`, which must run on the same task. No-op when the
+    span is not in flight. See `_safe_detach` for how this scope is used to
+    classify the cross-task detach it enables.
+
+    Args:
+        key: Correlation key of an in-flight span (the key it was stashed under).
+    """
+    entry = _in_flight_spans.get(key)
+    if entry is None:
+        return
+    span = entry[0]
+    token = otel_context.attach(trace.set_span_in_context(span))
+    _reattached_tokens[key] = (token, _current_task())
+
+
+def release_reattached_span(key: str) -> None:
+    """Release a reattached span from a matching `reattach_span()` call.
+
+    Must run on the same task that called `reattach_span()`. No-op when no token
+    is stored.
+
+    Args:
+        key: Correlation key from the matching `reattach_span()` call.
+    """
+    entry = _reattached_tokens.pop(key, None)
+    if entry is not None:
+        token, _ = entry
+        otel_context.detach(token)
+
+
 def finish_streaming_span(
     streaming_id: str,
     *,