fix(agentex): end SSE subscriptions on terminal task; stop deleting shared stream

Task event SSE subscriptions had no terminal condition. stream_task_events ran
a `while True` loop whose only exits were client disconnect (CancelledError) and
fatal errors. Once a task finished its producer stopped writing, but the reader
kept blocking on the topic forever — issuing an XREAD every couple of seconds
and pinning one connection from the shared per-process Redis pool. Because the
stream key carries a sliding TTL, a finished task's key eventually disappears
while readers keep blocking on it, turning each subscription into a permanent
zombie. Accumulated zombies exhaust the pool, which is shared with the readiness
probe, so /readyz fails while the dependency-free /healthz keeps returning 200 —
the pod goes Unready and is never restarted.

Termination:
- Primary is event-driven. Every terminal transition funnels through
  task_service.transition_status (and delete via update_task), which XADDs a
  task_updated event carrying the new status onto the task's own stream. The
  reader already delivers those events, so it returns as soon as it forwards one
  whose task is in a terminal status — no DB lookup and no ordering race, since
  the terminal event is the last message produced.
- Fallback runs on the keepalive-ping cadence for the cases the event path
  cannot see: a client that connects after the task already finished, a producer
  that dies without emitting a terminal event, or a task_updated with no task
  payload. It ends the stream when the task is terminal or when a topic that had
  existed is reclaimed. stream_ever_existed is seeded from the tail snapshot so a
  mid-flight subscriber correctly arms the vanished-topic check.

Shared stream:
- Drop the per-subscriber cleanup_stream in finally. The topic is keyed only by
  task id and shared by every viewer, so deleting it on one subscriber's exit
  tore the stream out from under the others. The sliding TTL already reclaims it.

Also adds stream_exists to the stream port + Redis adapter (EXISTS) for the
vanished-topic fallback, and two integration regression tests: the stream ends
on its own once the task is terminal, and a single subscriber disconnecting does
not delete the shared topic.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: deepthi-rao-scale

agentex/src/adapters/streams/adapter_redis.py

@@ -229,6 +229,14 @@ async def read_messages(
             logger.error(f"Error reading from Redis stream {topic}: {e}")
             raise
 
+    async def stream_exists(self, topic: str) -> bool:
+        """Whether the Redis stream key exists (EXISTS)."""
+        try:
+            return bool(await self.redis.exists(topic))
+        except Exception as e:
+            logger.error(f"Error checking existence of Redis stream {topic}: {e}")
+            raise
+
     async def cleanup_stream(self, topic: str) -> None:
         """
         Clean up a Redis stream.

agentex/src/adapters/streams/port.py

@@ -60,6 +60,14 @@ async def get_stream_tail_id(self, topic: str) -> str:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    async def stream_exists(self, topic: str) -> bool:
+        """
+        Report whether a stream topic currently exists, to tell a reclaimed
+        topic ("gone") apart from one that is still live.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     async def cleanup_stream(self, topic: str) -> None:
         """

agentex/src/domain/use_cases/streams_use_case.py

@@ -12,14 +12,28 @@
     TaskStreamConnectedEventEntity,
     TaskStreamErrorEventEntity,
     TaskStreamEventEntity,
+    TaskStreamTaskUpdatedEventEntity,
     convert_task_stream_event_to_entity,
 )
+from src.domain.entities.tasks import TaskStatus
 from src.domain.services.task_service import DAgentTaskService
 from src.utils.logging import make_logger
 from src.utils.stream_topics import get_task_event_stream_topic
 
 logger = make_logger(__name__)
 
+# Statuses after which no more events are produced; a subscription must end.
+_TERMINAL_TASK_STATUSES = frozenset(
+    {
+        TaskStatus.CANCELED,
+        TaskStatus.COMPLETED,
+        TaskStatus.FAILED,
+        TaskStatus.TERMINATED,
+        TaskStatus.TIMED_OUT,
+        TaskStatus.DELETED,
+    }
+)
+
 
 class StreamsUseCase:
     def __init__(
@@ -119,6 +133,13 @@ async def stream_task_events(
         # 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
+        # A subscription must end once no further events can arrive, or it
+        # blocks on the topic forever, pinning a Redis connection (a "zombie").
+        # Primary signal is event-driven (a terminal task_updated below); the
+        # keepalive-cadence check is a fallback. stream_ever_existed gates the
+        # fallback's vanished-topic check; a concrete tail ID means the stream
+        # already had data at connect, so it counts as having existed.
+        stream_ever_existed = last_id != "0-0"
         try:
             # Application-level control loop
             while True:
@@ -132,21 +153,46 @@ async def stream_task_events(
                         # Update the last_id for the next iteration
                         last_id = new_id
                         message_count += 1
+                        # Topic exists; a later disappearance now means "gone",
+                        # not "not yet created".
+                        stream_ever_existed = True
                         # 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()
+                        # A terminal task_updated is always the last event
+                        # produced, so nothing follows it — end the stream 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
 
-                    # If we didn't get any messages, add a small pause
-                    # to prevent tight loops and send keepalive ping if needed
+                    # No messages this cycle: on each keepalive interval, run the
+                    # fallback done-check (for finishes the event loop never
+                    # sees: late connect, silent producer death, a payload-less
+                    # task_updated) and, if still live, send a keepalive ping.
                     if message_count == 0:
                         current_time = asyncio.get_running_loop().time()
                         if current_time - last_message_time >= ping_interval:
+                            if await self._stream_is_finished(
+                                task_id, stream_topic, stream_ever_existed
+                            ):
+                                logger.info(
+                                    f"Ending SSE stream for task {task_id}: "
+                                    "fallback check found it finished"
+                                )
+                                break
                             yield ":ping\n\n"
                             last_message_time = current_time
                         await asyncio.sleep(0.1)
@@ -190,8 +236,42 @@ async def stream_task_events(
             )
             yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n"
         finally:
+            # No cleanup_stream here: the topic is shared by all viewers of the
+            # task, so deleting it on one subscriber's exit would break the
+            # others. The sliding TTL reclaims it instead.
             logger.info(f"SSE stream for task {task_id} has ended")
-            await self.cleanup_stream(stream_topic)
+
+    async def _stream_is_finished(
+        self,
+        task_id: str,
+        stream_topic: str,
+        stream_ever_existed: bool,
+    ) -> bool:
+        """
+        Fallback terminal check. True if the task is terminal, or the topic
+        was reclaimed after having existed (covers hard-deleted tasks whose
+        lookup fails). A not-yet-written topic is not treated as finished. A
+        transient lookup failure keeps the stream open.
+        """
+        try:
+            task = await self.task_service.get_task(id=task_id)
+        except Exception as e:
+            logger.warning(
+                f"Terminal-state lookup failed for task {task_id}; "
+                f"keeping stream open: {e}"
+            )
+            task = None
+        # Normal finish: the DB says the task is done, so no more XADDs.
+        if task is not None and task.status in _TERMINAL_TASK_STATUSES:
+            return True
+        # Vanished topic: the key existed before but its sliding TTL reclaimed
+        # it, so a reader is now blocking on a key that can never return data
+        # (also covers a hard-deleted task whose status lookup failed above).
+        if stream_ever_existed and not await self.stream_repository.stream_exists(
+            stream_topic
+        ):
+            return True
+        return False
 
 
 DStreamsUseCase = Annotated[StreamsUseCase, Depends(StreamsUseCase)]

agentex/tests/integration/test_task_stream.py

@@ -588,9 +588,7 @@ async def test_event_xadded_after_connected_is_delivered(
 
             # Emit a delta while the generator is suspended at the yield.
             sentinel = "after-connected-sentinel"
-            await repo.send_data(
-                stream_topic, {"type": "error", "message": sentinel}
-            )
+            await repo.send_data(stream_topic, {"type": "error", "message": sentinel})
 
             # The delta must be delivered; a silent stream means it was dropped.
             received = False
@@ -656,3 +654,93 @@ async def collect_stream_data():
         )
 
         print(f"✅ Stream sent {ping_count} keepalive pings during idle period")
+
+    async def test_stream_ends_when_task_reaches_terminal_status(
+        self, test_agent_and_task, tasks_use_case, streams_use_case
+    ):
+        """
+        Regression for the zombie-subscription leak: the stream must end on its
+        own once the task is terminal (no client-side cancellation), instead of
+        blocking on the topic forever and pinning a Redis connection.
+
+        Completing via TasksUseCase emits a terminal task_updated event, so this
+        exercises the event-driven termination path.
+        """
+        _agent, task = test_agent_and_task
+
+        async def drain_until_end():
+            # Returns normally only if the generator terminates by itself.
+            async for _event_data in streams_use_case.stream_task_events(
+                task_id=task.id
+            ):
+                pass
+
+        reader_task = asyncio.create_task(drain_until_end())
+
+        # Let the stream connect and enter its loop, then finish the task.
+        await asyncio.sleep(0.2)
+        completed = await tasks_use_case.complete_task(id=task.id)
+        assert completed.status == TaskStatus.COMPLETED
+
+        # The generator should now end by itself.
+        try:
+            await asyncio.wait_for(reader_task, timeout=10)
+        except TimeoutError as err:
+            reader_task.cancel()
+            try:
+                await reader_task
+            except asyncio.CancelledError:
+                pass
+            raise AssertionError(
+                "SSE stream did not terminate after the task reached a terminal "
+                "status — the subscription is a zombie and will pin a Redis "
+                "connection forever."
+            ) from err
+
+        print("✅ Stream ends on its own once the task is terminal")
+
+    async def test_cleanup_does_not_delete_shared_stream_on_disconnect(
+        self, test_agent_and_task, streams_use_case
+    ):
+        """
+        Regression for the shared-stream deletion bug: the topic is shared by
+        every viewer of a task, so one subscriber disconnecting must not delete
+        it (the old per-subscriber cleanup_stream did). The sliding TTL handles
+        reclamation instead.
+        """
+        from src.utils.stream_topics import get_task_event_stream_topic
+
+        _agent, task = test_agent_and_task
+        stream_topic = get_task_event_stream_topic(task_id=task.id)
+        repo = streams_use_case.stream_repository
+
+        # Ensure the topic exists.
+        await repo.send_data(stream_topic, {"type": "error", "message": "seed"})
+        assert await repo.stream_exists(stream_topic), "precondition: topic exists"
+
+        # Run one subscriber briefly, then disconnect it (simulating one viewer
+        # of a task that other viewers are still watching).
+        async def brief_reader():
+            try:
+                async for _event_data in streams_use_case.stream_task_events(
+                    task_id=task.id
+                ):
+                    pass
+            except asyncio.CancelledError:
+                pass
+
+        reader_task = asyncio.create_task(brief_reader())
+        await asyncio.sleep(0.3)
+        reader_task.cancel()
+        try:
+            await reader_task
+        except asyncio.CancelledError:
+            pass
+
+        # The shared topic must survive one subscriber leaving.
+        assert await repo.stream_exists(stream_topic), (
+            "Topic was deleted when a single subscriber disconnected — other "
+            "live viewers of this task would lose their stream."
+        )
+
+        print("✅ Shared stream survives a single subscriber disconnecting")