fix(agentex): terminate SSE task streams deterministically; stop leaking Redis connections

Task event SSE subscriptions had no terminal condition: stream_task_events ran a
`while True` loop that only exited on client disconnect or a fatal error. Once a
task finished, its producer stopped writing but the reader kept blocking on the
Redis stream (XREAD) forever, pinning a connection from the shared per-process
pool. Because stream keys carry a sliding TTL, a finished task's key eventually
expires while readers keep blocking on it — a permanent zombie. Accumulated
zombies exhaust the pool, which is shared with the readiness probe, so /readyz
fails while the dependency-free /healthz stays 200: the pod goes Unready and is
never restarted.

Termination now has three independent, deterministic checks:
- Connect-time: read task status once at connect (after snapshotting the cursor,
  so a racing terminal event still lands after the cursor). If already terminal,
  replay buffered events and end — handles late connects.
- In-stream: end on a terminal task_updated event (the last event a task emits).
- Periodic authoritative recheck: at the top of every loop iteration, on an
  interval, re-read task status and end if terminal. Runs busy or idle and even
  after a read failure/backoff, so a dropped/lost terminal event or a failing
  read cannot keep the stream open. A hard-deleted task (ItemDoesNotExist) is
  treated as terminal; other lookup errors fall through to transient retry.

Also:
- AgentTaskService.fail_task now publishes task_updated via update_task — it was
  the only terminal write that didn't emit, which could strand a live viewer.
- Drop the per-subscriber cleanup_stream in `finally`: the topic is shared by all
  viewers, so deleting it on one exit broke the others. The sliding TTL reclaims.
- Canonical NON_TERMINAL/TERMINAL task-status sets on the TaskStatus entity,
  referenced by both the status state machine and SSE termination.

Adds integration regression tests: event-driven termination, late-connect
termination, a running task surviving a reclaimed key, and shared-stream
survival on disconnect.

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

Author: deepthi-rao-scale

agentex/src/domain/entities/tasks.py

@@ -29,6 +29,13 @@ class TaskStatus(str, Enum):
     DELETED = "DELETED"
 
 
+# Canonical status partition (state machine + SSE termination).
+# Non-terminal: RUNNING or INTERRUPTED (resumable); terminal is the rest.
+# New statuses are terminal unless added to the non-terminal set.
+NON_TERMINAL_TASK_STATUSES = frozenset({TaskStatus.RUNNING, TaskStatus.INTERRUPTED})
+TERMINAL_TASK_STATUSES = frozenset(TaskStatus) - NON_TERMINAL_TASK_STATUSES
+
+
 class TaskEntity(BaseModel):
     id: str = Field(
         ...,

agentex/src/domain/services/task_service.py

@@ -166,7 +166,9 @@ async def forward_task_to_acp(
     async def fail_task(self, task: TaskEntity, reason: str) -> None:
         task.status = TaskStatus.FAILED
         task.status_reason = reason
-        await self.task_repository.update(task)
+        # Publish task_updated so streaming viewers see the failure and end.
+        # Every terminal write must emit; SSE termination relies on it.
+        await self.update_task(task)
 
     async def get_task(
         self,
@@ -374,7 +376,11 @@ async def cancel_task(
             new_status=TaskStatus.CANCELED,
             status_reason="Task canceled by user",
         )
-        return updated if updated is not None else await self.task_repository.get(id=task.id)
+        return (
+            updated
+            if updated is not None
+            else await self.task_repository.get(id=task.id)
+        )
 
     async def interrupt_task(
         self, agent: AgentEntity, task: TaskEntity, acp_url: str

agentex/src/domain/use_cases/streams_use_case.py

@@ -5,15 +5,18 @@
 from fastapi import Depends
 from pydantic import ValidationError
 
+from src.adapters.crud_store.exceptions import ItemDoesNotExist
 from src.adapters.streams.adapter_redis import DRedisStreamRepository
 from src.api.schemas.task_stream_events import TaskStreamEvent
 from src.config.dependencies import DEnvironmentVariables
 from src.domain.entities.task_stream_events import (
     TaskStreamConnectedEventEntity,
     TaskStreamErrorEventEntity,
     TaskStreamEventEntity,
+    TaskStreamTaskUpdatedEventEntity,
     convert_task_stream_event_to_entity,
 )
+from src.domain.entities.tasks import TERMINAL_TASK_STATUSES
 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
@@ -101,15 +104,21 @@ 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.
+        # 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
@@ -119,10 +128,34 @@ 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
+        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, 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
@@ -136,16 +169,25 @@ async def stream_task_events(
                         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
 
-                    # If we didn't get any messages, add a small pause
-                    # to prevent tight loops and send keepalive ping if needed
+                    # Idle: send keepalive ping so proxies don't reap us.
                     if message_count == 0:
-                        current_time = asyncio.get_running_loop().time()
                         if current_time - last_message_time >= ping_interval:
                             yield ":ping\n\n"
                             last_message_time = current_time
@@ -190,8 +232,8 @@ async def stream_task_events(
             )
             yield f"data: {TaskStreamErrorEventEntity(type='error', message=str(e)).model_dump_json()}\n\n"
         finally:
+            # Don't delete the shared topic; the TTL reclaims it.
             logger.info(f"SSE stream for task {task_id} has ended")
-            await self.cleanup_stream(stream_topic)
 
 
 DStreamsUseCase = Annotated[StreamsUseCase, Depends(StreamsUseCase)]

agentex/src/domain/use_cases/tasks_use_case.py

@@ -3,7 +3,12 @@
 from fastapi import Depends
 
 from src.adapters.crud_store.exceptions import ItemDoesNotExist
-from src.domain.entities.tasks import TaskEntity, TaskRelationships, TaskStatus
+from src.domain.entities.tasks import (
+    NON_TERMINAL_TASK_STATUSES,
+    TaskEntity,
+    TaskRelationships,
+    TaskStatus,
+)
 from src.domain.exceptions import ClientError
 from src.domain.services.task_service import DAgentTaskService
 from src.utils.logging import make_logger
@@ -132,10 +137,8 @@ async def update_mutable_fields_on_task(
 
         return task_entity
 
-    # Non-terminal statuses a task can be transitioned to a terminal status from.
-    # RUNNING is the normal case; INTERRUPTED is also valid so an interrupted
-    # (paused, still-continuable) task can still be canceled/completed/etc later.
-    _TERMINAL_TRANSITION_SOURCES = (TaskStatus.RUNNING, TaskStatus.INTERRUPTED)
+    # Statuses a task can transition to terminal from (the non-terminal set).
+    _TERMINAL_TRANSITION_SOURCES = NON_TERMINAL_TASK_STATUSES
 
     async def _transition_to_terminal(
         self,

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,198 @@ 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 bool(await repo.redis.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 bool(await repo.redis.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")
+
+    async def test_stream_ends_on_late_connect_to_terminal_task(
+        self, test_agent_and_task, tasks_use_case, streams_use_case
+    ):
+        """
+        Connect-time termination path: a viewer that connects *after* the task
+        has already reached a terminal state never receives a terminal
+        task_updated event via the read loop — the event is behind the snapshot
+        cursor, so the in-loop check can't fire. The authoritative connect-time
+        check must end the stream (replay buffered events, then return) instead
+        of blocking on the topic forever.
+        """
+        _agent, task = test_agent_and_task
+
+        # Finish the task BEFORE anyone subscribes. The terminal task_updated is
+        # XADDed to the stream now, so a later subscriber snapshots past it and
+        # the read loop never surfaces it.
+        completed = await tasks_use_case.complete_task(id=task.id)
+        assert completed.status == TaskStatus.COMPLETED
+
+        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())
+
+        # The connect-time check should end it near-immediately.
+        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 to an already-terminal task did not self-close — the "
+                "connect-time terminal check is not firing, so late-connect "
+                "subscriptions leak as zombies."
+            ) from err
+
+        print("✅ Stream self-closes when connecting to an already-terminal task")
+
+    async def test_running_task_stream_survives_reclaimed_key(
+        self, test_agent_and_task, streams_use_case
+    ):
+        """
+        A still-RUNNING task whose idle stream key is reclaimed by the sliding
+        TTL must NOT have its stream closed — a later XADD recreates the key.
+        Termination is driven by the terminal task_updated event (and the
+        connect-time terminal check), never by a vanished key, so a live task's
+        stream stays open regardless of key reclamation.
+        """
+        from src.utils.stream_topics import get_task_event_stream_topic
+
+        # Short keepalive so several idle cycles elapse within the test window.
+        streams_use_case.environment_variables.SSE_KEEPALIVE_PING_INTERVAL = 1
+
+        _agent, task = test_agent_and_task  # created in RUNNING state
+        stream_topic = get_task_event_stream_topic(task_id=task.id)
+        repo = streams_use_case.stream_repository
+
+        # Seed the topic, subscribe, then reclaim the key mid-stream.
+        await repo.send_data(stream_topic, {"type": "error", "message": "seed"})
+        assert bool(await repo.redis.exists(stream_topic)), "precondition: topic exists"
+
+        async def 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(reader())
+
+        # Let the reader connect and go idle, then simulate the sliding TTL
+        # reclaiming the key while the task is still RUNNING.
+        await asyncio.sleep(0.5)
+        await repo.cleanup_stream(stream_topic)
+        assert not bool(await repo.redis.exists(stream_topic)), (
+            "precondition: key reclaimed"
+        )
+
+        # Give the fallback several cycles to (wrongly) close the stream.
+        await asyncio.sleep(4)
+
+        try:
+            assert not reader_task.done(), (
+                "SSE stream for a still-RUNNING task was closed after its idle key "
+                "was reclaimed — the fallback must not treat a reclaimed key as "
+                "terminal when the task is confirmed non-terminal."
+            )
+        finally:
+            reader_task.cancel()
+            try:
+                await reader_task
+            except asyncio.CancelledError:
+                pass
+
+        print("✅ Running task's stream stays open when its idle key is reclaimed")