fix(tasks): guarantee stream open/close pairing; tag duration by outcome

Move record_stream_opened() inside the try/finally and gate record_stream_closed
on an `opened` flag, so a failure between marking the stream open and entering
the try can no longer leave the active gauge permanently over-counted. Also tag
the duration histogram with `outcome` (matching the closed counter) so stream
lifetime can be sliced by completed / client_disconnect / error on both the OTel
and StatsD paths.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: cyntwang99

agentex/src/domain/use_cases/streams_use_case.py

@@ -108,15 +108,20 @@ async def stream_task_events(
 
         stream_topic = get_task_event_stream_topic(task_id=task_id)
 
-        # A stream is "open" from here on. Pair this with exactly one
-        # record_stream_closed (in the finally below) so the active-stream gauge
-        # stays balanced even if setup — the tail snapshot or the first yield —
-        # raises. Placed after task resolution so a bad task_name never counts
+        # 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.
-        record_stream_opened()
         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
@@ -234,10 +239,13 @@ async def stream_task_events(
             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")
-            record_stream_closed(
-                outcome,
-                asyncio.get_running_loop().time() - stream_start_time,
-            )
+            # 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)
 
 

agentex/src/utils/stream_metrics.py

@@ -133,7 +133,8 @@ def record_stream_opened() -> None:
 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, and lowers the active gauge.
+    outcome), records the stream lifetime (also tagged by outcome), and lowers
+    the active gauge.
 
     Args:
         outcome: One of "completed", "client_disconnect", "error".
@@ -148,14 +149,18 @@ def record_stream_closed(outcome: StreamOutcome, duration_seconds: float) -> Non
         if _closed_counter is not None:
             _closed_counter.add(1, {"outcome": outcome})
         if _duration_histogram is not None:
-            _duration_histogram.record(duration_seconds)
+            _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)
+            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).
     except Exception:

agentex/tests/unit/utils/test_stream_metrics.py

@@ -87,9 +87,11 @@ def test_record_stream_closed_emits_statsd_when_enabled():
         "agentex.task_stream.closed",
         tags=["outcome:client_disconnect"],
     )
-    # Duration is reported to StatsD in milliseconds.
+    # Duration is reported to StatsD in milliseconds, tagged by outcome.
     mock_statsd.histogram.assert_called_once_with(
-        "agentex.task_stream.duration", 3250.0
+        "agentex.task_stream.duration",
+        3250.0,
+        tags=["outcome:client_disconnect"],
     )
     # Concurrency ("active") is OTel-only, so closing emits no StatsD gauge.
     mock_statsd.gauge.assert_not_called()
@@ -132,7 +134,7 @@ def test_closed_records_otel_instruments_with_outcome():
 
     assert opened.calls == [(1, None)]
     assert closed.calls == [(1, {"outcome": "error"})]
-    assert duration.calls == [(4.0, None)]
+    assert duration.calls == [(4.0, {"outcome": "error"})]
     # +1 on open, -1 on close nets to a balanced active gauge.
     assert active.calls == [(1, None), (-1, None)]