From 0d2be849e91a72e4b01bc0f39da87881a684350e Mon Sep 17 00:00:00 2001 From: Cynthia Wang Date: Fri, 31 Jul 2026 15:47:46 -0700 Subject: [PATCH 1/3] fix(observability): isolate SSE stream spans and continue ingress traceparent Each task-event stream now runs under its own OpenTelemetry span, parented only on the inbound W3C traceparent (or a fresh root when absent) and never the ambient context. The SSE body is pumped by the ASGI server after the request handler returns, so the ambient context could still carry a prior request's span; inheriting it made a stream's logs and child spans resolve to an unrelated, long-lived trace (cross-request context bleed). The span carries open/first-event/close lifecycle events and the task.id / stream.outcome / disconnect.reason attributes, and stays attached for the generator's lifetime so log lines correlate to the correct trace. Co-Authored-By: Claude Opus 4.7 --- agentex/src/api/routes/tasks.py | 16 +- .../src/domain/use_cases/streams_use_case.py | 376 +++++++++++------- .../test_streams_use_case_tracing.py | 206 ++++++++++ 3 files changed, 456 insertions(+), 142 deletions(-) create mode 100644 agentex/tests/unit/use_cases/test_streams_use_case_tracing.py 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 e8f14da3..cb0cd148 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 @@ -23,6 +29,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__( @@ -90,151 +101,238 @@ 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. + + 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. """ - 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) - # Cursor before status read: catches a racing terminal event. - last_id = await self.stream_repository.get_stream_tail_id(stream_topic) - task = await self.task_service.get_task(id=task_id) - # Send initial connection data - yield f"data: {TaskStreamConnectedEventEntity(type='connected', taskId=task_id).model_dump_json()}\n\n" - # Already terminal: replay buffered events and end (late connect). - if task.status in TERMINAL_TASK_STATUSES: - async for _id, data in self.read_messages(topic=stream_topic, last_id="0"): - yield f"data: {data.model_dump_json()}\n\n" - await asyncio.sleep(0.02) - logger.info( - f"Ending SSE stream for task {task_id}: already terminal at connect" - ) - return - - last_message_time = asyncio.get_running_loop().time() - ping_interval = float( - self.environment_variables.SSE_KEEPALIVE_PING_INTERVAL - ) # Configurable keepalive ping interval - # Track consecutive read failures so we can back off and avoid a - # tight error loop. When the Redis pool is exhausted, every connected - # client's read fails on each cycle; without backoff this turns into a - # log-ingestion firehose (one failure per client per cycle, ~once/sec). - consecutive_errors = 0 - last_status_check = last_message_time + # 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) + ) + + # Lifecycle outcome, finalized on the span in ``finally``. Defaults assume + # a clean server-side end; each termination path narrows the reason. + stream_outcome = "completed" + disconnect_reason = "completed" + 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: - # Application-level control loop - 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. - 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. - logger.info( - f"Ending SSE stream for task {task_id}: " - "task no longer exists" - ) - return - if task.status in TERMINAL_TASK_STATUSES: - logger.info( - f"Ending SSE stream for task {task_id}: " - "terminal on status recheck" - ) - return - - # Process yielded messages one by one - message_generator = self.read_messages( - topic=stream_topic, last_id=last_id - ) - message_count = 0 - async for new_id, data in message_generator: - # Update the last_id for the next iteration - last_id = new_id - message_count += 1 - # Send the data to the client - data_str = f"data: {data.model_dump_json()}\n\n" - yield data_str - last_message_time = asyncio.get_running_loop().time() - # Terminal event is the last one — end here. - if ( - isinstance(data, TaskStreamTaskUpdatedEventEntity) - and data.task is not None - and data.task.status in TERMINAL_TASK_STATUSES - ): - logger.info( - f"Ending SSE stream for task {task_id}: received " - "a terminal task_updated event" - ) - return - await asyncio.sleep(0.02) - - # A read cycle completed without raising — the stream is - # healthy again, so reset the backoff/error counter. - 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. - if message_count == 0: - now = asyncio.get_running_loop().time() - if now - last_message_time >= ping_interval: - yield ":ping\n\n" - last_message_time = now - await asyncio.sleep(0.1) - else: - # Small pause between batches - await asyncio.sleep(0.02) - except asyncio.CancelledError: - # Client disconnected, exit the loop - logger.info( - f"Client disconnected from SSE stream for task {task_id}" - ) - raise - except Exception as e: - consecutive_errors += 1 - # Always log the full traceback — nothing is swallowed. - # Volume is controlled two ways instead of by dropping - # diagnostics: structured JSON logging keeps each traceback - # to a single log entry (see utils.logging), and the - # exponential backoff below caps how often a sustained - # failure can repeat. The failure counter gives context on - # how long a stream has been erroring. - logger.error( - f"Error processing events for task {task_id} " - f"(failure #{consecutive_errors}): {e}", - exc_info=True, - ) - yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n" - # Exponential backoff (capped) so a sustained failure (e.g. - # Redis pool exhaustion) doesn't spin a tight per-client - # loop hammering Redis and flooding logs. - backoff = min(2.0 ** min(consecutive_errors - 1, 5), 30.0) - await asyncio.sleep(backoff) - - except asyncio.CancelledError: - # Just exit the generator on cancellation - logger.info(f"Client disconnected from SSE stream for task {task_id}") - pass - except Exception as 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" + 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) + # Cursor before status read: catches a racing terminal event. + last_id = await self.stream_repository.get_stream_tail_id(stream_topic) + task = await self.task_service.get_task(id=task_id) + # Send initial connection data + yield f"data: {TaskStreamConnectedEventEntity(type='connected', taskId=task_id).model_dump_json()}\n\n" + # Already terminal: replay buffered events and end (late connect). + if task.status in TERMINAL_TASK_STATUSES: + 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" + ) + return + + last_message_time = asyncio.get_running_loop().time() + ping_interval = float( + self.environment_variables.SSE_KEEPALIVE_PING_INTERVAL + ) # Configurable keepalive ping interval + # Track consecutive read failures so we can back off and avoid a + # tight error loop. When the Redis pool is exhausted, every connected + # client's read fails on each cycle; without backoff this turns into a + # log-ingestion firehose (one failure per client per cycle, ~once/sec). + consecutive_errors = 0 + last_status_check = last_message_time + try: + # Application-level control loop + 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. + 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. + stream_outcome = "task_gone" + 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" + ) + return + + # Process yielded messages one by one + message_generator = self.read_messages( + topic=stream_topic, last_id=last_id + ) + message_count = 0 + async for new_id, data in message_generator: + # 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 + last_message_time = asyncio.get_running_loop().time() + # Terminal event is the last one — end here. + if ( + isinstance(data, TaskStreamTaskUpdatedEventEntity) + 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" + ) + return + await asyncio.sleep(0.02) + + # A read cycle completed without raising — the stream is + # healthy again, so reset the backoff/error counter. + 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. + if message_count == 0: + now = asyncio.get_running_loop().time() + if now - last_message_time >= ping_interval: + yield ":ping\n\n" + last_message_time = now + await asyncio.sleep(0.1) + else: + # Small pause between batches + await asyncio.sleep(0.02) + except asyncio.CancelledError: + # Client disconnected, exit the loop + stream_outcome = "client_disconnect" + disconnect_reason = "client_disconnect" + logger.info( + f"Client disconnected from SSE stream for task {task_id}" + ) + raise + except Exception as e: + consecutive_errors += 1 + # Always log the full traceback — nothing is swallowed. + # Volume is controlled two ways instead of by dropping + # diagnostics: structured JSON logging keeps each traceback + # to a single log entry (see utils.logging), and the + # exponential backoff below caps how often a sustained + # failure can repeat. The failure counter gives context on + # how long a stream has been erroring. + logger.error( + f"Error processing events for task {task_id} " + f"(failure #{consecutive_errors}): {e}", + exc_info=True, + ) + yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n" + # Exponential backoff (capped) so a sustained failure (e.g. + # Redis pool exhaustion) doesn't spin a tight per-client + # loop hammering Redis and flooding logs. + backoff = min(2.0 ** min(consecutive_errors - 1, 5), 30.0) + await asyncio.sleep(backoff) + + except asyncio.CancelledError: + # Just exit the generator on cancellation + stream_outcome = "client_disconnect" + disconnect_reason = "client_disconnect" + logger.info(f"Client disconnected from SSE stream for task {task_id}") + pass + except Exception as e: + stream_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: + # Classify an outcome that no handler already set from the exception + # still in flight. Two uncaught cases reach here: Starlette aborting + # the body with GeneratorExit on client disconnect, and an exception + # from the pre-loop setup (task resolution, initial reads) — neither + # passes through the loop's handlers, so without this they'd be + # mislabeled "completed". + in_flight = sys.exc_info()[1] + if in_flight is not None and stream_outcome == "completed": + if isinstance(in_flight, GeneratorExit | asyncio.CancelledError): + stream_outcome = "client_disconnect" + disconnect_reason = "client_disconnect" + else: + stream_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", stream_outcome) + span.set_attribute("disconnect.reason", disconnect_reason) + span.add_event( + "close", + { + "stream.outcome": stream_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") 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..18790670 --- /dev/null +++ b/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py @@ -0,0 +1,206 @@ +"""Tracing tests for the SSE task-event stream. + +These lock in the AGX1-617 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), + ) + + +@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")) + + # 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 From 6424ac0c27c2dea492479eeb5182ff4b853e34cc Mon Sep 17 00:00:00 2001 From: Cynthia Wang Date: Mon, 3 Aug 2026 09:53:07 -0700 Subject: [PATCH 2/3] fix(observability): emit SSE error frame on stream setup failure Merge the stream generator's two try blocks into one so the fatal except handlers cover the setup phase (task resolution, stream tail read, initial status read, terminal replay) as well as the read loop. Previously a setup failure escaped the generator instead of producing the established SSE error frame; the streaming response wrapper then raised StreamResponseError, so clients saw a broken stream. Now such a failure yields a `data: {"type":"error",...}` frame, ends the stream cleanly, and marks the span stream.outcome=error / disconnect.reason with the exception type. Co-Authored-By: Claude Opus 4.7 --- .../src/domain/use_cases/streams_use_case.py | 238 +++++++++--------- .../test_streams_use_case_tracing.py | 26 ++ 2 files changed, 144 insertions(+), 120 deletions(-) diff --git a/agentex/src/domain/use_cases/streams_use_case.py b/agentex/src/domain/use_cases/streams_use_case.py index cb0cd148..151d8cb2 100644 --- a/agentex/src/domain/use_cases/streams_use_case.py +++ b/agentex/src/domain/use_cases/streams_use_case.py @@ -187,131 +187,129 @@ def _record_first_event() -> None: # log-ingestion firehose (one failure per client per cycle, ~once/sec). consecutive_errors = 0 last_status_check = last_message_time - try: - # Application-level control loop - 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. - 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. - stream_outcome = "task_gone" - 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" - ) - return + # Application-level control loop + 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. + 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. + stream_outcome = "task_gone" + 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" + ) + return - # Process yielded messages one by one - message_generator = self.read_messages( - topic=stream_topic, last_id=last_id - ) - message_count = 0 - async for new_id, data in message_generator: - # 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 - last_message_time = asyncio.get_running_loop().time() - # Terminal event is the last one — end here. - if ( - isinstance(data, TaskStreamTaskUpdatedEventEntity) - 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" - ) - return - await asyncio.sleep(0.02) + # Process yielded messages one by one + message_generator = self.read_messages( + topic=stream_topic, last_id=last_id + ) + message_count = 0 + async for new_id, data in message_generator: + # 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 + last_message_time = asyncio.get_running_loop().time() + # Terminal event is the last one — end here. + if ( + isinstance(data, TaskStreamTaskUpdatedEventEntity) + 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" + ) + return + await asyncio.sleep(0.02) - # A read cycle completed without raising — the stream is - # healthy again, so reset the backoff/error counter. - consecutive_errors = 0 + # A read cycle completed without raising — the stream is + # healthy again, so reset the backoff/error counter. + 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. - if message_count == 0: - now = asyncio.get_running_loop().time() - if now - last_message_time >= ping_interval: - yield ":ping\n\n" - last_message_time = now - await asyncio.sleep(0.1) - else: - # Small pause between batches - await asyncio.sleep(0.02) - except asyncio.CancelledError: - # Client disconnected, exit the loop - stream_outcome = "client_disconnect" - disconnect_reason = "client_disconnect" - logger.info( - f"Client disconnected from SSE stream for task {task_id}" - ) - raise - except Exception as e: - consecutive_errors += 1 - # Always log the full traceback — nothing is swallowed. - # Volume is controlled two ways instead of by dropping - # diagnostics: structured JSON logging keeps each traceback - # to a single log entry (see utils.logging), and the - # exponential backoff below caps how often a sustained - # failure can repeat. The failure counter gives context on - # how long a stream has been erroring. - logger.error( - f"Error processing events for task {task_id} " - f"(failure #{consecutive_errors}): {e}", - exc_info=True, - ) - yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n" - # Exponential backoff (capped) so a sustained failure (e.g. - # Redis pool exhaustion) doesn't spin a tight per-client - # loop hammering Redis and flooding logs. - backoff = min(2.0 ** min(consecutive_errors - 1, 5), 30.0) - await asyncio.sleep(backoff) + # 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. + if message_count == 0: + now = asyncio.get_running_loop().time() + if now - last_message_time >= ping_interval: + yield ":ping\n\n" + last_message_time = now + await asyncio.sleep(0.1) + else: + # Small pause between batches + await asyncio.sleep(0.02) + except asyncio.CancelledError: + # Client disconnected, exit the loop + stream_outcome = "client_disconnect" + disconnect_reason = "client_disconnect" + logger.info( + f"Client disconnected from SSE stream for task {task_id}" + ) + raise + except Exception as e: + consecutive_errors += 1 + # Always log the full traceback — nothing is swallowed. + # Volume is controlled two ways instead of by dropping + # diagnostics: structured JSON logging keeps each traceback + # to a single log entry (see utils.logging), and the + # exponential backoff below caps how often a sustained + # failure can repeat. The failure counter gives context on + # how long a stream has been erroring. + logger.error( + f"Error processing events for task {task_id} " + f"(failure #{consecutive_errors}): {e}", + exc_info=True, + ) + yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n" + # Exponential backoff (capped) so a sustained failure (e.g. + # Redis pool exhaustion) doesn't spin a tight per-client + # loop hammering Redis and flooding logs. + backoff = min(2.0 ** min(consecutive_errors - 1, 5), 30.0) + await asyncio.sleep(backoff) - except asyncio.CancelledError: - # Just exit the generator on cancellation - stream_outcome = "client_disconnect" - disconnect_reason = "client_disconnect" - logger.info(f"Client disconnected from SSE stream for task {task_id}") - pass - except Exception as e: - stream_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" + except asyncio.CancelledError: + # Just exit the generator on cancellation + stream_outcome = "client_disconnect" + disconnect_reason = "client_disconnect" + logger.info(f"Client disconnected from SSE stream for task {task_id}") + pass + except Exception as e: + stream_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: - # Classify an outcome that no handler already set from the exception - # still in flight. Two uncaught cases reach here: Starlette aborting - # the body with GeneratorExit on client disconnect, and an exception - # from the pre-loop setup (task resolution, initial reads) — neither - # passes through the loop's handlers, so without this they'd be - # mislabeled "completed". + # 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 stream_outcome == "completed": if isinstance(in_flight, GeneratorExit | asyncio.CancelledError): 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 index 18790670..63461300 100644 --- a/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py +++ b/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py @@ -204,3 +204,29 @@ async def test_span_ends_exactly_once_and_detaches_context(self, 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), + ) + 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" From 54166f5c89516604f117bf5a17389be121427c59 Mon Sep 17 00:00:00 2001 From: Cynthia Wang Date: Wed, 5 Aug 2026 10:37:15 -0700 Subject: [PATCH 3/3] test(observability): address SSE stream tracing review comments Remove an internal tracker ID from the tracing test docstring (this repo is public) and assert exactly one stream span is finished in the context-detach test so it matches its name. Co-Authored-By: Claude Opus 4.7 --- .../tests/unit/use_cases/test_streams_use_case_tracing.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 index abbab8b2..2a61ed09 100644 --- a/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py +++ b/agentex/tests/unit/use_cases/test_streams_use_case_tracing.py @@ -1,6 +1,6 @@ """Tracing tests for the SSE task-event stream. -These lock in the AGX1-617 fixes: +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 @@ -204,6 +204,9 @@ async def test_span_ends_exactly_once_and_detaches_context(self, span_exporter): 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