fix(observability): isolate SSE stream spans and continue ingress traceparent

agentex/src/api/routes/tasks.py

@@ -1,7 +1,7 @@
 import json
 from typing import Annotated, Any
 
-from fastapi import APIRouter, HTTPException, Query
+from fastapi import APIRouter, HTTPException, Query, Request
 from fastapi.responses import StreamingResponse
 
 from src.adapters.temporal.adapter_temporal import DTemporalAdapter
@@ -345,13 +345,18 @@ async def timeout_task(
 async def stream_task_events(
     task_id: DAuthorizedId(AgentexResourceType.task, AuthorizedOperationType.read),
     stream_use_case: DStreamsUseCase,
+    request: Request,
 ) -> StreamingResponse:
     """
     Streams task events using Server-Sent Events (SSE).
     """
 
     return StreamingResponse(
-        stream_use_case.stream_task_events(task_id=task_id),
+        # Pass request headers so the stream span continues an inbound W3C
+        # traceparent (ingress edge) instead of starting an unrelated trace.
+        stream_use_case.stream_task_events(
+            task_id=task_id, carrier=dict(request.headers)
+        ),
         media_type="text/event-stream",
         headers={
             "Cache-Control": "no-cache",
@@ -369,13 +374,18 @@ async def stream_task_events(
 async def stream_task_events_by_name(
     task_name: DAuthorizedName(AgentexResourceType.task, AuthorizedOperationType.read),
     stream_use_case: DStreamsUseCase,
+    request: Request,
 ) -> StreamingResponse:
     """
     Streams task events using Server-Sent Events (SSE) by task name.
     """
 
     return StreamingResponse(
-        stream_use_case.stream_task_events(task_name=task_name),
+        # Pass request headers so the stream span continues an inbound W3C
+        # traceparent (ingress edge) instead of starting an unrelated trace.
+        stream_use_case.stream_task_events(
+            task_name=task_name, carrier=dict(request.headers)
+        ),
         media_type="text/event-stream",
         headers={
             "Cache-Control": "no-cache",

agentex/src/domain/use_cases/streams_use_case.py

@@ -1,8 +1,14 @@
 import asyncio
-from collections.abc import AsyncIterator
+import sys
+from collections.abc import AsyncIterator, Mapping
 from typing import Annotated
 
 from fastapi import Depends
+from opentelemetry import context as otel_context
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.propagate import extract
+from opentelemetry.trace import SpanKind, Status, StatusCode
 from pydantic import ValidationError
 
 from src.adapters.crud_store.exceptions import ItemDoesNotExist
@@ -29,6 +35,11 @@
 
 logger = make_logger(__name__)
 
+# ProxyTracer: resolves to the process TracerProvider lazily at span-creation
+# time, so this is safe to bind at import even when tracing is configured later
+# (and a no-op when no provider is installed).
+_TRACER = trace.get_tracer("agentex.task_stream")
+
 
 class StreamsUseCase:
     def __init__(
@@ -96,33 +107,71 @@ async def stream_task_events(
         self,
         task_id: str | None = None,
         task_name: str | None = None,
+        carrier: Mapping[str, str] | None = None,
     ) -> AsyncIterator[str]:
         """
         Async generator for streaming task message updates as SSE data strings.
         Sends keepalive pings to maintain long-lived connections.
-        """
-        task_id = task_id
-        if not task_id:
-            if not task_name:
-                raise ValueError("Either task_id or task_name must be provided")
-
-            task = await self.task_service.get_task(name=task_name)
-            task_id = task.id
 
-        stream_topic = get_task_event_stream_topic(task_id=task_id)
+        Each call runs under its own OpenTelemetry span whose parent is taken
+        *only* from the inbound W3C ``traceparent`` (``carrier``) or, absent one,
+        a fresh root — never the ambient context. The SSE body is pumped by the
+        ASGI server long after the request handler returned, so the ambient
+        context can still carry a previous request's span; inheriting it made a
+        stream's logs and child spans resolve to an unrelated, long-lived trace
+        (cross-request context bleed). Anchoring to an isolated context and
+        attaching it for the stream's lifetime keeps every stream self-contained.
+        """
+        # Parent the stream span on the ingress traceparent alone. The empty
+        # ``Context()`` base is the isolation: ``extract`` otherwise falls back to
+        # the current (possibly stale) context, which is exactly the bleed. With
+        # no inbound traceparent this yields a fresh root; with one, a child.
+        parent_context = extract(dict(carrier) if carrier else {}, context=Context())
+        span = _TRACER.start_span(
+            "stream task events",
+            context=parent_context,
+            kind=SpanKind.SERVER,
+        )
+        # Attach the span as current for the stream's lifetime so log lines
+        # correlate to this trace (fixing the "otelTraceID resolves to an
+        # unrelated trace" symptom) and any downstream spans nest under it.
+        context_token = otel_context.attach(
+            trace.set_span_in_context(span, parent_context)
+        )
 
-        # Capture the timing/outcome state the finally needs *before* marking the
-        # stream open, then flip ``opened`` as the first statement inside the try.
-        # This keeps record_stream_opened paired with exactly one
-        # record_stream_closed: the open lives inside the try, so any failure
-        # after it still routes through the finally and rebalances the active
-        # gauge — closing the narrow window that opening before the try left
-        # exposed. Placed after task resolution so a bad task_name never counts
-        # as an opened stream.
+        # Lifecycle state finalized in ``finally``. ``outcome`` is the coarse label
+        # shared by the close metric and the span's ``stream.outcome`` attribute;
+        # ``disconnect_reason`` is the finer, span-only detail. ``stream_start_time``
+        # and ``opened`` let the finally emit exactly one balanced close metric.
         stream_start_time = asyncio.get_running_loop().time()
         outcome: StreamOutcome = "completed"
+        disconnect_reason = "completed"
         opened = False
+        first_event_recorded = False
+
+        def _record_first_event() -> None:
+            # Time-to-first-event marker: the first real task event delivered
+            # (the synthetic "connected" frame and keepalive pings don't count).
+            nonlocal first_event_recorded
+            if not first_event_recorded:
+                first_event_recorded = True
+                span.add_event("first-event")
+
+        span.add_event("open")
         try:
+            # Resolve task_name -> id inside the try so a failure yields the SSE
+            # error frame (and marks the span errored) instead of escaping into a
+            # broken stream. Mark the stream opened only after resolution so a bad
+            # task_name never counts as an opened stream.
+            if not task_id:
+                if not task_name:
+                    raise ValueError("Either task_id or task_name must be provided")
+
+                task = await self.task_service.get_task(name=task_name)
+                task_id = task.id
+            span.set_attribute("task.id", task_id)
+
+            stream_topic = get_task_event_stream_topic(task_id=task_id)
             record_stream_opened()
             opened = True
             # Snapshot the read cursor BEFORE yielding "connected". "connected"
@@ -142,8 +191,10 @@ async def stream_task_events(
                 async for _id, data in self.read_messages(
                     topic=stream_topic, last_id="0"
                 ):
+                    _record_first_event()
                     yield f"data: {data.model_dump_json()}\n\n"
                     await asyncio.sleep(0.02)
+                disconnect_reason = "already_terminal"
                 logger.info(
                     f"Ending SSE stream for task {task_id}: already terminal at connect"
                 )
@@ -175,22 +226,26 @@ async def stream_task_events(
             while True:
                 try:
                     # Authoritative status recheck on an interval. Runs at the
-                    # TOP of every iteration — even after a read failure/backoff —
-                    # so a terminal task ends even if its event publish was lost
-                    # or Redis reads keep erroring.
+                    # TOP of every iteration — even after a read failure/backoff
+                    # — so a terminal task ends even if its event publish was
+                    # lost or Redis reads keep erroring.
                     current_time = asyncio.get_running_loop().time()
                     if current_time - last_status_check >= ping_interval:
                         last_status_check = current_time
                         try:
                             task = await self.task_service.get_task(id=task_id)
                         except ItemDoesNotExist:
-                            # Row permanently gone (e.g. retention) — end, don't retry.
+                            # Row permanently gone (e.g. retention) — end. Coarse
+                            # outcome stays "completed"; the fine reason records
+                            # that the underlying row was deleted.
+                            disconnect_reason = "task_deleted"
                             logger.info(
                                 f"Ending SSE stream for task {task_id}: "
                                 "task no longer exists"
                             )
                             return
                         if task.status in TERMINAL_TASK_STATUSES:
+                            disconnect_reason = "terminal_status"
                             logger.info(
                                 f"Ending SSE stream for task {task_id}: "
                                 "terminal on status recheck"
@@ -206,6 +261,7 @@ async def stream_task_events(
                         # Update the last_id for the next iteration
                         last_id = new_id
                         message_count += 1
+                        _record_first_event()
                         # Send the data to the client
                         data_str = f"data: {data.model_dump_json()}\n\n"
                         yield data_str
@@ -219,6 +275,7 @@ async def stream_task_events(
                             and data.task is not None
                             and data.task.status in TERMINAL_TASK_STATUSES
                         ):
+                            disconnect_reason = "terminal_event"
                             logger.info(
                                 f"Ending SSE stream for task {task_id}: received "
                                 "a terminal task_updated event"
@@ -231,8 +288,8 @@ async def stream_task_events(
                     consecutive_errors = 0
 
                     # Idle: send keepalive ping so proxies don't reap us. Use a
-                    # fresh timestamp — the read above blocks up to timeout_ms, so
-                    # the loop-top current_time would be stale for ping timing.
+                    # fresh timestamp — the read above blocks up to timeout_ms,
+                    # so the loop-top current_time would be stale for ping timing.
                     if message_count == 0:
                         now = asyncio.get_running_loop().time()
                         # No data event pushed for the stall window: count the
@@ -251,6 +308,8 @@ async def stream_task_events(
                         await asyncio.sleep(0.02)
                 except asyncio.CancelledError:
                     # Client disconnected, exit the loop
+                    outcome = "client_disconnect"
+                    disconnect_reason = "client_disconnect"
                     logger.info(
                         f"Client disconnected from SSE stream for task {task_id}"
                     )
@@ -279,14 +338,44 @@ async def stream_task_events(
         except asyncio.CancelledError:
             # Just exit the generator on cancellation
             outcome = "client_disconnect"
+            disconnect_reason = "client_disconnect"
             logger.info(f"Client disconnected from SSE stream for task {task_id}")
         except Exception as e:
             outcome = "error"
+            disconnect_reason = type(e).__name__
+            span.record_exception(e)
+            span.set_status(Status(StatusCode.ERROR, str(e)))
             logger.error(
                 f"Fatal error in SSE stream for task {task_id}: {e}", exc_info=True
             )
             yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n"
         finally:
+            # Safety net for an exception that bypassed the except clauses above
+            # and is still in flight — chiefly Starlette aborting the SSE body
+            # with GeneratorExit on client disconnect, which is a BaseException
+            # that `except Exception` doesn't see. Only reclassify when no
+            # handler already set an outcome.
+            in_flight = sys.exc_info()[1]
+            if in_flight is not None and outcome == "completed":
+                if isinstance(in_flight, GeneratorExit | asyncio.CancelledError):
+                    outcome = "client_disconnect"
+                    disconnect_reason = "client_disconnect"
+                else:
+                    outcome = "error"
+                    disconnect_reason = type(in_flight).__name__
+                    span.record_exception(in_flight)
+                    span.set_status(Status(StatusCode.ERROR, str(in_flight)))
+            span.set_attribute("stream.outcome", outcome)
+            span.set_attribute("disconnect.reason", disconnect_reason)
+            span.add_event(
+                "close",
+                {
+                    "stream.outcome": outcome,
+                    "disconnect.reason": disconnect_reason,
+                },
+            )
+            span.end()
+            otel_context.detach(context_token)
             # Don't delete the shared topic; the TTL reclaims it.
             logger.info(f"SSE stream for task {task_id} has ended")
             # Only close what we actually opened, so the active gauge is never