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

## What was broken

When you open a task in the UI, the backend keeps a live stream open to
push updates. That stream had **no way to end itself** — it only stopped
if the browser disconnected.

So when a task finished, the backend just kept listening forever,
holding one Redis connection each time. These pile up until the pool is
empty, at which point the pod can't serve requests (and its health check
fails, so it sits stuck until a manual restart).

A second bug: when one viewer left, the code deleted the task's
**shared** event stream — yanking it out from under everyone else
watching the same task.

## The fix

- **End the stream when the task is done.** The backend already gets a
`task_updated` event when a task finishes, so the stream now closes as
soon as it sees a terminal status.
- **Fallback for the rare cases** the event can't cover (viewer connects
after the task already finished, or the producer dies silently): a
lightweight check on the existing keepalive interval closes the stream
then too.
- **Stop deleting the shared stream** on one viewer's exit — the Redis
TTL already cleans it up on its own.

## Testing

Two integration regression tests (stream self-ends on terminal task; one
viewer disconnecting doesn't delete the shared stream).

**Verified live (A/B against `main`)** — same steps both times: create a
RUNNING task → open `GET /tasks/{id}/stream` → complete the task, while
watching Redis `blocked_clients`.

| | on task completion | Redis connection |
|---|---|---|
| **`main`** | stream keeps `XREAD`-looping; kept sending `:ping` 15s+
after completion and only ended when the client was force-disconnected |
`blocked_clients` stayed `1` — released only on client disconnect
(**zombie**) |
| **this branch** | stream delivers the terminal `task_updated` and
returns immediately (`Ending SSE stream …: received a terminal
task_updated event`) | `blocked_clients` `1 → 0` instantly, no client
disconnect needed |

On `main` the reader logged `Reading messages from Redis stream …` every
~2s indefinitely after the task was done; on this branch it stops the
moment the terminal event arrives. Same reproduction the incident
described, and the connection is released instead of pinned.

> Note: the unit/integration suite runs in CI; the tutorial-agent
integration jobs were red due to a pre-existing missing-`pytest-asyncio`
issue unrelated to this change (fixed separately in
scale-agentex-python).

Related: the UI-side half of this leak is #383.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- greptile_comment -->

<h3>Greptile Summary</h3>

This PR updates task SSE subscription lifecycle handling.
- Ends subscriptions when a terminal update is received, authoritative
task state becomes terminal, or the task disappears.
- Replays buffered events and closes when connecting to an
already-terminal task.
- Publishes failure updates and centralizes terminal/non-terminal status
classification.
- Preserves shared Redis streams when individual subscribers disconnect.
- Adds integration coverage for terminal shutdown, late connections,
reclaimed keys, transient lookups, and shared-stream retention.

<details><summary><h3>Confidence Score: 4/5</h3></summary>

The PR should not merge until recoverable FAILED tasks no longer
terminate subscriptions that must observe their subsequent RUNNING
updates.

The stream returns upon receiving any FAILED task update, while the task
service intentionally supports forwarding that same FAILED task again
and publishing a RUNNING recovery update; viewers whose subscriptions
closed on FAILED cannot receive that recovery or subsequent events.

**Files Needing Attention:** agentex/src/domain/entities/tasks.py,
agentex/src/domain/use_cases/streams_use_case.py,
agentex/src/domain/services/task_service.py
</details>

<details><summary><h3>Important Files Changed</h3></summary>

| Filename | Overview |
|----------|----------|
| agentex/src/domain/entities/tasks.py | Introduces canonical terminal
and non-terminal status sets shared by state transitions and SSE
termination. |
| agentex/src/domain/services/task_service.py | Publishes FAILED updates
through the normal update path, while existing retry behavior can still
restore FAILED tasks to RUNNING. |
| agentex/src/domain/use_cases/streams_use_case.py | Adds event-driven
and authoritative fallback termination and stops deleting the shared
Redis topic, but the previously reported recoverable-FAILED lifecycle
conflict remains. |
| agentex/src/domain/use_cases/tasks_use_case.py | Reuses the canonical
non-terminal status set for transitions to terminal states. |
| agentex/tests/integration/test_task_stream.py | Adds regression
coverage for terminal shutdown, late subscribers, vanished tasks,
transient lookups, reclaimed keys, and shared-stream preservation. |

</details>

<details><summary><h3>Sequence Diagram</h3></summary>

```mermaid
sequenceDiagram
    participant Client
    participant SSE as StreamsUseCase
    participant DB as Task Store
    participant Redis
    Client->>SSE: Subscribe to task stream
    SSE->>Redis: Snapshot stream cursor
    SSE->>DB: Read authoritative task state
    alt Task already terminal
        SSE->>Redis: Replay buffered events
        SSE-->>Client: Events, then close
    else Task is non-terminal
        loop Until terminal, missing, or disconnect
            SSE->>Redis: Read new events
            alt Terminal task_updated received
                SSE-->>Client: Terminal event
                SSE-->>Client: Close stream
            else Status-check interval elapsed
                SSE->>DB: Recheck task
                alt Terminal or missing
                    SSE-->>Client: Close stream
                end
            end
        end
    end
    Note over SSE,Redis: Subscriber exit does not delete the shared stream
```
</details>

<sub>Reviews (14): Last reviewed commit: ["fix(agentex): terminate SSE
task
streams..."](85dfe6b)
| [Re-trigger
Greptile](https://app.greptile.com/api/retrigger?id=48824816)</sub>

<!-- /greptile_comment -->

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,