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,19 +169,31 @@ 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. Use a
+                    # fresh timestamp — the read above blocks up to timeout_ms, so
+                    # the loop-top current_time would be stale for ping timing.
                     if message_count == 0:
-                        current_time = asyncio.get_running_loop().time()
-                        if current_time - last_message_time >= ping_interval:
+                        now = asyncio.get_running_loop().time()
+                        if now - last_message_time >= ping_interval:
                             yield ":ping\n\n"
-                            last_message_time = current_time
+                            last_message_time = now
                         await asyncio.sleep(0.1)
                     else:
                         # Small pause between batches
@@ -190,8 +235,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,