Skip to content
8 changes: 8 additions & 0 deletions agentex/src/config/environment_variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ class EnvVarKeys(str, Enum):
HTTPX_POOL_TIMEOUT = "HTTPX_POOL_TIMEOUT"
HTTPX_STREAMING_READ_TIMEOUT = "HTTPX_STREAMING_READ_TIMEOUT"
SSE_KEEPALIVE_PING_INTERVAL = "SSE_KEEPALIVE_PING_INTERVAL"
SSE_STREAM_STALL_THRESHOLD_SECONDS = "SSE_STREAM_STALL_THRESHOLD_SECONDS"
AGENTEX_SERVER_TASK_QUEUE = "AGENTEX_SERVER_TASK_QUEUE"
ENABLE_HEALTH_CHECK_WORKFLOW = "ENABLE_HEALTH_CHECK_WORKFLOW"
ENABLE_AGENT_RUN_SCHEDULES = "ENABLE_AGENT_RUN_SCHEDULES"
Expand Down Expand Up @@ -150,6 +151,10 @@ class EnvironmentVariables(BaseModel):
300.0 # HTTPX streaming read timeout in seconds (5 minutes)
)
SSE_KEEPALIVE_PING_INTERVAL: int = 15 # SSE keepalive ping interval in seconds
# An open stream is counted as stalled once this many seconds pass with no
# data event pushed to the client. Kept above the keepalive interval so that
# keepalive pings (which are not data events) don't mask a real stall.
SSE_STREAM_STALL_THRESHOLD_SECONDS: int = 30
AGENTEX_SERVER_TASK_QUEUE: str | None = None
ENABLE_HEALTH_CHECK_WORKFLOW: bool = False
# Gates the agent run schedules API. Off by default; enabled in development.
Expand Down Expand Up @@ -245,6 +250,9 @@ def refresh(cls, force_refresh: bool = False) -> EnvironmentVariables | None:
SSE_KEEPALIVE_PING_INTERVAL=int(
os.environ.get(EnvVarKeys.SSE_KEEPALIVE_PING_INTERVAL, "15")
),
SSE_STREAM_STALL_THRESHOLD_SECONDS=int(
os.environ.get(EnvVarKeys.SSE_STREAM_STALL_THRESHOLD_SECONDS, "30")
),
AGENTEX_SERVER_TASK_QUEUE=os.environ.get(
EnvVarKeys.AGENTEX_SERVER_TASK_QUEUE
),
Expand Down
95 changes: 75 additions & 20 deletions agentex/src/domain/use_cases/streams_use_case.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@
)
from src.domain.services.task_service import DAgentTaskService
from src.utils.logging import make_logger
from src.utils.stream_metrics import (
StreamOutcome,
record_stream_closed,
record_stream_opened,
record_stream_stall,
)
from src.utils.stream_topics import get_task_event_stream_topic

logger = make_logger(__name__)
Expand Down Expand Up @@ -101,25 +107,53 @@ async def stream_task_events(
task_id = task.id

stream_topic = get_task_event_stream_topic(task_id=task_id)
# Snapshot the read cursor BEFORE yielding "connected". "connected" is
# the client's cue to send its message, which makes the agent start
# XADD-ing deltas. Snapshotting after the yield lets a congested relay
# fall behind far enough that those deltas land before the snapshot and
# are never read. Snapshotting first resolves to "0-0" (stream is empty
# until the client sends), so we read from the beginning.
last_id = await self.stream_repository.get_stream_tail_id(stream_topic)
# Send initial connection data
yield f"data: {TaskStreamConnectedEventEntity(type='connected', taskId=task_id).model_dump_json()}\n\n"
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

# 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.
stream_start_time = asyncio.get_running_loop().time()
outcome: StreamOutcome = "completed"
opened = False
try:
record_stream_opened()
opened = True
# Snapshot the read cursor BEFORE yielding "connected". "connected"
# is the client's cue to send its message, which makes the agent
# start XADD-ing deltas. Snapshotting after the yield lets a
# congested relay fall behind far enough that those deltas land
# before the snapshot and are never read. Snapshotting first resolves
# to "0-0" (stream is empty until the client sends), so we read from
# the beginning.
last_id = await self.stream_repository.get_stream_tail_id(stream_topic)
# Send initial connection data
yield f"data: {TaskStreamConnectedEventEntity(type='connected', taskId=task_id).model_dump_json()}\n\n"
now = asyncio.get_running_loop().time()
# last_message_time drives the keepalive ping and is reset by pings;
# last_event_time tracks only real data pushes (a ping is not a data
# event) and drives stall detection, so a persistently quiet stream
# still trips the stall signal even while keepalive pings flow.
last_message_time = now
last_event_time = now
ping_interval = float(
self.environment_variables.SSE_KEEPALIVE_PING_INTERVAL
) # Configurable keepalive ping interval
stall_threshold = float(
self.environment_variables.SSE_STREAM_STALL_THRESHOLD_SECONDS
)
# Whether the current quiet spell has already been counted, so one
# stall episode increments the counter once (measuring stall onsets)
# rather than once per idle cycle. Reset when a data event flows.
stalled = False
# 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
# Application-level control loop
while True:
try:
Expand All @@ -135,7 +169,10 @@ async def stream_task_events(
# 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()
now = asyncio.get_running_loop().time()
last_message_time = now
last_event_time = now
stalled = False
await asyncio.sleep(0.02)

# A read cycle completed without raising — the stream is
Expand All @@ -146,6 +183,16 @@ async def stream_task_events(
# to prevent tight loops and send keepalive ping if needed
if message_count == 0:
current_time = asyncio.get_running_loop().time()
# No data event pushed for the stall window: count the
# onset once. Keepalive pings deliberately do not reset
# last_event_time, so a persistently quiet stream is
# visible even though the connection stays alive.
if (
current_time - last_event_time >= stall_threshold
and not stalled
):
stalled = True
record_stream_stall()
if current_time - last_message_time >= ping_interval:
yield ":ping\n\n"
last_message_time = current_time
Expand Down Expand Up @@ -182,15 +229,23 @@ async def stream_task_events(

except asyncio.CancelledError:
# Just exit the generator on cancellation
outcome = "client_disconnect"
logger.info(f"Client disconnected from SSE stream for task {task_id}")
pass
except Exception as e:
outcome = "error"
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:
logger.info(f"SSE stream for task {task_id} has ended")
# Only close what we actually opened, so the active gauge is never
# decremented for a stream that never incremented it.
if opened:
record_stream_closed(
outcome,
asyncio.get_running_loop().time() - stream_start_time,
)
await self.cleanup_stream(stream_topic)


Expand Down
187 changes: 187 additions & 0 deletions agentex/src/utils/stream_metrics.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,187 @@
"""
Metrics instrumentation for the task-event SSE stream lifecycle.

Mirrors the dual-emit pattern in ``src/utils/cache_metrics.py``:

- When an OTLP endpoint is configured (``OTEL_EXPORTER_OTLP_ENDPOINT``), the
instruments are recorded through the OpenTelemetry SDK.
- When the Datadog Agent is reachable (``DD_AGENT_HOST``), the same events are
emitted as StatsD metrics.
- When neither is configured, every function here is a cheap no-op.

**Why this exists (see AGX1-616/AGX1-618):** HTTP request-level RED for
``GET /tasks/{task_id}/stream`` is already covered by auto-instrumentation — the
route records into ``http_server_request_duration_seconds`` at connection close.
But a single SSE request produces exactly one duration sample, which says nothing
about what happened *during* the stream's life: how many were concurrently open,
whether a stream ended normally vs. by client disconnect vs. by error, or whether
a stream went quiet (no events pushed) for a stretch. Those are the stream-
lifecycle signals request-level instrumentation cannot express, and they are what
this module adds.

Instrument names are deliberately dotted (``agentex.task_stream.*``); the OTLP →
Prometheus translation flattens dots to underscores, appends ``_total`` to
monotonic counters, and appends the unit (``s`` → ``_seconds``) to the histogram,
so the series land in Mimir as ``agentex_task_stream_opened_total``,
``agentex_task_stream_closed_total``, ``agentex_task_stream_duration_seconds``,
``agentex_task_stream_active`` and ``agentex_task_stream_stall_total``. These
hand-instrumented series carry only bounded attributes (``outcome``) and no
``http_route`` label — group by ``outcome``, not ``http_route``.
"""

from __future__ import annotations

import os
from typing import TYPE_CHECKING, Literal

from datadog import statsd

from src.utils.logging import make_logger
from src.utils.otel_metrics import get_meter

if TYPE_CHECKING:
from opentelemetry.metrics import Counter, Histogram, UpDownCounter

logger = make_logger(__name__)

# StatsD is only emitted if the Datadog Agent host is configured.
_STATSD_ENABLED = bool(os.environ.get("DD_AGENT_HOST"))

# How a stream ended. "completed" = the generator returned normally;
# "client_disconnect" = the client went away (asyncio.CancelledError);
# "error" = the stream loop raised a fatal, non-recoverable exception. HTTP
# status 200 at connection close cannot tell these three apart, which is why the
# distinction lives here as a bounded label rather than on the request metric.
StreamOutcome = Literal["completed", "client_disconnect", "error"]

# Lazily-created OTel instruments (created once, on first use).
_opened_counter: Counter | None = None
_closed_counter: Counter | None = None
_duration_histogram: Histogram | None = None
_active_updown: UpDownCounter | None = None
_stall_counter: Counter | None = None
_instruments_initialized = False


def _ensure_instruments() -> None:
"""Create the OTel instruments on first use. No-op if OTel is not configured."""
global _opened_counter, _closed_counter, _duration_histogram
global _active_updown, _stall_counter, _instruments_initialized

if _instruments_initialized:
return
_instruments_initialized = True

meter = get_meter("agentex.task_stream")
if meter is None:
# OTel not configured; OTel path stays disabled. StatsD may still emit.
return

_opened_counter = meter.create_counter(
name="agentex.task_stream.opened",
description="Task-event SSE streams opened",
unit="{stream}",
)
_closed_counter = meter.create_counter(
name="agentex.task_stream.closed",
description="Task-event SSE streams closed, tagged by outcome",
unit="{stream}",
)
_duration_histogram = meter.create_histogram(
name="agentex.task_stream.duration",
description="Task-event SSE stream lifetime",
unit="s",
)
_active_updown = meter.create_up_down_counter(
name="agentex.task_stream.active",
description="Currently open task-event SSE streams",
unit="{stream}",
)
_stall_counter = meter.create_counter(
name="agentex.task_stream.stall",
description="Task-event SSE streams that went idle (no event pushed) past the stall threshold",
unit="{stream}",
)


def record_stream_opened() -> None:
"""
Record a task-event stream opening: bumps the opened counter and the active gauge.

Pair every call with exactly one ``record_stream_closed`` so the active gauge
stays balanced. Never raises: see ``record_stream_closed``.
"""
try:
_ensure_instruments()
Comment thread
greptile-apps[bot] marked this conversation as resolved.

if _opened_counter is not None:
_opened_counter.add(1)
if _active_updown is not None:
_active_updown.add(1)

if _STATSD_ENABLED:
statsd.increment("agentex.task_stream.opened")
# No StatsD "active" gauge: concurrency is OTel-only. DogStatsD
# gauges are last-write-wins per flush, so N workers each emitting a
# process-local count would make the gauge flap between workers
# rather than sum to the true concurrency. The OTel UpDownCounter
# sums correctly across per-instance series at query time.
except Exception:
logger.debug("Failed to emit agentex.task_stream.opened metric", exc_info=True)


def record_stream_closed(outcome: StreamOutcome, duration_seconds: float) -> None:
"""
Record a task-event stream closing: bumps the closed counter (tagged by
outcome), records the stream lifetime (also tagged by outcome), and lowers
the active gauge.

Args:
outcome: One of "completed", "client_disconnect", "error".
duration_seconds: How long the stream was open.

Never raises: emission failures (e.g. a StatsD UDP socket error or an OTel
SDK fault) are swallowed so instrumentation can never disrupt the SSE path.
"""
try:
_ensure_instruments()

if _closed_counter is not None:
_closed_counter.add(1, {"outcome": outcome})
if _duration_histogram is not None:
_duration_histogram.record(duration_seconds, {"outcome": outcome})
if _active_updown is not None:
_active_updown.add(-1)

if _STATSD_ENABLED:
statsd.increment("agentex.task_stream.closed", tags=[f"outcome:{outcome}"])
# Datadog histograms conventionally take milliseconds for durations.
statsd.histogram(
"agentex.task_stream.duration",
duration_seconds * 1000,
tags=[f"outcome:{outcome}"],
)
# No StatsD "active" gauge — concurrency is OTel-only; a per-worker
# DogStatsD gauge would flap rather than sum (see record_stream_opened).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given that these are new metrics, and in the final state we want otel to emit to datadog instead of emitting to dd directly, i think we can remove the double emission here.

except Exception:
logger.debug("Failed to emit agentex.task_stream.closed metric", exc_info=True)


def record_stream_stall() -> None:
"""
Record that an open stream went idle — no event pushed to the client for
longer than the configured stall threshold. Emit once per stall episode
(the caller de-dupes) so the counter measures stall onsets, not idle cycles.

Never raises: see ``record_stream_closed``.
"""
try:
_ensure_instruments()

if _stall_counter is not None:
_stall_counter.add(1)

if _STATSD_ENABLED:
statsd.increment("agentex.task_stream.stall")
except Exception:
logger.debug("Failed to emit agentex.task_stream.stall metric", exc_info=True)
Loading
Loading