diff --git a/agentex/src/api/routes/tasks.py b/agentex/src/api/routes/tasks.py index 6c9dab91..6d3607c7 100644 --- a/agentex/src/api/routes/tasks.py +++ b/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", diff --git a/agentex/src/domain/use_cases/streams_use_case.py b/agentex/src/domain/use_cases/streams_use_case.py index 559c3e0e..32b292fa 100644 --- a/agentex/src/domain/use_cases/streams_use_case.py +++ b/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 diff --git a/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py b/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py new file mode 100644 index 00000000..2a61ed09 --- /dev/null +++ b/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py @@ -0,0 +1,241 @@ +"""Tracing tests for the SSE task-event stream. + +These lock in the SSE trace context bleed fixes: + +1. Each stream runs under its own span, isolated from the ambient OTel context — + a stream must never nest under a leftover span from an unrelated request + (the cross-request "context bleed" that made a /stream log's trace resolve to + a multi-hour, thousands-of-spans trace). +2. When the caller supplies a W3C ``traceparent`` (ingress edge), the stream + span continues that trace as a child rather than starting a new root. +3. The span carries the ``open`` / ``first-event`` / ``close`` lifecycle events + and the ``task.id`` / ``stream.outcome`` / ``disconnect.reason`` attributes. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from types import SimpleNamespace + +import pytest +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from src.domain.entities.tasks import TaskStatus +from src.domain.use_cases import streams_use_case as streams_module +from src.domain.use_cases.streams_use_case import StreamsUseCase + +# A well-formed W3C traceparent (version-traceid-spanid-flags), sampled. +_INBOUND_TRACE_ID = 0x4BF92F3577B34DA6A3CE929D0E0E4736 +_INBOUND_PARENT_SPAN_ID = 0x00F067AA0BA902B7 +_INBOUND_TRACEPARENT = f"00-{_INBOUND_TRACE_ID:032x}-{_INBOUND_PARENT_SPAN_ID:016x}-01" + +_SPAN_NAME = "stream task events" + + +class _FakeStreamRepository: + """Minimal stream repo: a fixed tail id and a fixed buffered replay.""" + + def __init__(self, buffered: list[tuple[str, dict]] | None = None): + self._buffered = buffered or [] + + async def get_stream_tail_id(self, topic: str) -> str: + return "0-0" + + async def read_messages( + self, topic: str, last_id: str, timeout_ms: int = 2000, count: int = 10 + ) -> AsyncIterator[tuple[str, dict]]: + for message_id, obj in self._buffered: + yield message_id, obj + + +class _FakeTaskService: + """Returns a single task for both id and name lookups.""" + + def __init__(self, task: SimpleNamespace): + self._task = task + + async def get_task(self, id=None, name=None) -> SimpleNamespace: + return self._task + + +def _make_use_case( + *, + status: TaskStatus, + buffered: list[tuple[str, dict]] | None = None, + task_id: str = "task-123", +) -> StreamsUseCase: + task = SimpleNamespace(id=task_id, status=status) + return StreamsUseCase( + stream_repository=_FakeStreamRepository(buffered), + task_service=_FakeTaskService(task), + environment_variables=SimpleNamespace( + SSE_KEEPALIVE_PING_INTERVAL=15, + SSE_STREAM_STALL_THRESHOLD_SECONDS=30, + ), + ) + + +@pytest.fixture +def span_exporter(monkeypatch) -> InMemorySpanExporter: + """Route the module's tracer to an in-memory exporter for assertions.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(streams_module, "_TRACER", provider.get_tracer("test")) + # Stash the provider so tests can mint a "stale" ambient span from it. + exporter._test_provider = provider # type: ignore[attr-defined] + return exporter + + +def _only_stream_span(exporter: InMemorySpanExporter): + spans = [s for s in exporter.get_finished_spans() if s.name == _SPAN_NAME] + assert len(spans) == 1, f"expected exactly one stream span, got {len(spans)}" + return spans[0] + + +async def _drain(gen) -> None: + async for _ in gen: + pass + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestStreamTracing: + async def test_stream_span_is_isolated_root_when_no_traceparent( + self, span_exporter + ): + # A stale span is active in the ambient context, mimicking a previous + # request whose context leaked. The stream must NOT nest under it. + provider = span_exporter._test_provider # type: ignore[attr-defined] + stale = provider.get_tracer("stale").start_span("POST /unrelated") + token = otel_context.attach(trace.set_span_in_context(stale)) + try: + uc = _make_use_case(status=TaskStatus.COMPLETED) + await _drain(uc.stream_task_events(task_id="task-123")) + finally: + otel_context.detach(token) + stale.end() + + span = _only_stream_span(span_exporter) + # No parent → a brand-new root, on a different trace than the stale span. + assert span.parent is None + assert span.context.trace_id != stale.get_span_context().trace_id + + async def test_stream_span_continues_inbound_traceparent(self, span_exporter): + uc = _make_use_case(status=TaskStatus.COMPLETED) + await _drain( + uc.stream_task_events( + task_id="task-123", + carrier={"traceparent": _INBOUND_TRACEPARENT}, + ) + ) + + span = _only_stream_span(span_exporter) + # Same trace as the caller, parented on the caller's span id. + assert span.context.trace_id == _INBOUND_TRACE_ID + assert span.parent is not None + assert span.parent.span_id == _INBOUND_PARENT_SPAN_ID + + async def test_inbound_traceparent_wins_over_ambient_context(self, span_exporter): + # Even with a (stale) ambient span active, the ingress traceparent — not + # the ambient context — must decide the parent. + provider = span_exporter._test_provider # type: ignore[attr-defined] + stale = provider.get_tracer("stale").start_span("POST /unrelated") + token = otel_context.attach(trace.set_span_in_context(stale)) + try: + uc = _make_use_case(status=TaskStatus.COMPLETED) + await _drain( + uc.stream_task_events( + task_id="task-123", + carrier={"traceparent": _INBOUND_TRACEPARENT}, + ) + ) + finally: + otel_context.detach(token) + stale.end() + + span = _only_stream_span(span_exporter) + assert span.context.trace_id == _INBOUND_TRACE_ID + assert span.context.trace_id != stale.get_span_context().trace_id + + async def test_lifecycle_events_and_attributes_on_clean_end(self, span_exporter): + # Already-terminal task replays one buffered event, then ends cleanly. + buffered = [("1-0", {"type": "error", "message": "buffered"})] + uc = _make_use_case(status=TaskStatus.COMPLETED, buffered=buffered) + await _drain(uc.stream_task_events(task_id="task-123")) + + span = _only_stream_span(span_exporter) + event_names = [e.name for e in span.events] + assert event_names == ["open", "first-event", "close"] + assert span.attributes["task.id"] == "task-123" + assert span.attributes["stream.outcome"] == "completed" + assert span.attributes["disconnect.reason"] == "already_terminal" + + async def test_first_event_absent_when_no_events_delivered(self, span_exporter): + # Terminal-at-connect with nothing buffered: open/close only. + uc = _make_use_case(status=TaskStatus.COMPLETED) + await _drain(uc.stream_task_events(task_id="task-123")) + + span = _only_stream_span(span_exporter) + assert [e.name for e in span.events] == ["open", "close"] + + async def test_client_disconnect_is_recorded_on_aclose(self, span_exporter): + # Non-terminal task: the generator suspends at the "connected" yield. + # Closing it (as Starlette does on client disconnect) raises GeneratorExit + # into the generator, which must be classified as a client disconnect. + uc = _make_use_case(status=TaskStatus.RUNNING) + gen = uc.stream_task_events(task_id="task-123") + first = await gen.__anext__() + assert "connected" in first + await gen.aclose() + + span = _only_stream_span(span_exporter) + assert span.attributes["stream.outcome"] == "client_disconnect" + assert span.attributes["disconnect.reason"] == "client_disconnect" + + async def test_span_ends_exactly_once_and_detaches_context(self, span_exporter): + # After the stream ends, the ambient context must be clean (the attached + # stream context was detached), so a later span is a fresh root. + uc = _make_use_case(status=TaskStatus.COMPLETED) + await _drain(uc.stream_task_events(task_id="task-123")) + + # Exactly one stream span was finished (ended once, not zero or twice). + _only_stream_span(span_exporter) + + # No stream span should be lingering as the current span. + current = trace.get_current_span() + assert current.get_span_context().trace_id == 0 # INVALID → nothing attached + + async def test_setup_failure_emits_error_frame_and_marks_span(self, span_exporter): + # A failure during setup (here, the initial task lookup) must produce the + # established SSE error frame rather than escaping the generator — which + # would surface to the client as a broken stream — and the span must be + # marked errored. + class _BoomTaskService: + async def get_task(self, id=None, name=None): + raise RuntimeError("boom") + + uc = StreamsUseCase( + stream_repository=_FakeStreamRepository(), + task_service=_BoomTaskService(), + environment_variables=SimpleNamespace( + SSE_KEEPALIVE_PING_INTERVAL=15, + SSE_STREAM_STALL_THRESHOLD_SECONDS=30, + ), + ) + frames = [chunk async for chunk in uc.stream_task_events(task_id="task-123")] + + # The generator completed normally, yielding a parseable SSE error frame. + assert frames, "expected an SSE error frame, got nothing" + assert frames[-1].startswith("data: ") + assert '"type":"error"' in frames[-1] + assert "boom" in frames[-1] + + span = _only_stream_span(span_exporter) + assert span.attributes["stream.outcome"] == "error" + assert span.attributes["disconnect.reason"] == "RuntimeError"