From 2eed3af10d68b38222af7383de739e2ac84bb9f8 Mon Sep 17 00:00:00 2001 From: Deepthi Rao Date: Wed, 29 Jul 2026 14:15:30 -0400 Subject: [PATCH 1/2] fix(agentex): end SSE subscriptions on terminal task; stop deleting shared stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- agentex/src/adapters/streams/adapter_redis.py | 8 ++ agentex/src/adapters/streams/port.py | 8 ++ .../src/domain/use_cases/streams_use_case.py | 86 ++++++++++++++++- agentex/tests/integration/test_task_stream.py | 94 ++++++++++++++++++- 4 files changed, 190 insertions(+), 6 deletions(-) diff --git a/agentex/src/adapters/streams/adapter_redis.py b/agentex/src/adapters/streams/adapter_redis.py index f1a20ea8..3894937e 100644 --- a/agentex/src/adapters/streams/adapter_redis.py +++ b/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. diff --git a/agentex/src/adapters/streams/port.py b/agentex/src/adapters/streams/port.py index 910c7a0e..78278b6b 100644 --- a/agentex/src/adapters/streams/port.py +++ b/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: """ diff --git a/agentex/src/domain/use_cases/streams_use_case.py b/agentex/src/domain/use_cases/streams_use_case.py index e25adffa..0c760fa9 100644 --- a/agentex/src/domain/use_cases/streams_use_case.py +++ b/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)] diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index 7da23edd..c65fd1b6 100644 --- a/agentex/tests/integration/test_task_stream.py +++ b/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") From d9d2fa60fb6fc0d794bf43d202ce574c6b6bb20f Mon Sep 17 00:00:00 2001 From: Deepthi Rao Date: Thu, 30 Jul 2026 12:23:38 -0400 Subject: [PATCH 2/2] docs: tighten the zombie-subscription comment in stream_task_events Co-Authored-By: Claude Opus 4.8 (1M context) --- agentex/src/domain/use_cases/streams_use_case.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/agentex/src/domain/use_cases/streams_use_case.py b/agentex/src/domain/use_cases/streams_use_case.py index 0c760fa9..13b69e8d 100644 --- a/agentex/src/domain/use_cases/streams_use_case.py +++ b/agentex/src/domain/use_cases/streams_use_case.py @@ -133,12 +133,11 @@ 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. + # 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