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..13b69e8d 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,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)] 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")