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

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,12 @@ 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
+        # End the sub once no more events can arrive, else XREAD blocks forever,
+        # pinning a Redis connection ("zombie"). Ends via a terminal task_updated
+        # (primary) or the keepalive fallback. A non-"0-0" tail means data existed
+        # at connect; stream_ever_existed gates the fallback so a not-yet-created
+        # topic isn't mistaken for a reclaimed one.
+        stream_ever_existed = last_id != "0-0"
         try:
             # Application-level control loop
             while True:
@@ -132,21 +152,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 +235,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")