From 85dfe6b7b124b4fa52fa932ae1b248d969a70c46 Mon Sep 17 00:00:00 2001 From: Deepthi Rao Date: Fri, 31 Jul 2026 11:58:22 -0400 Subject: [PATCH] fix(agentex): terminate SSE task streams deterministically; stop leaking Redis connections 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 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) --- agentex/src/domain/entities/tasks.py | 7 + agentex/src/domain/services/task_service.py | 10 +- .../src/domain/use_cases/streams_use_case.py | 69 ++++-- .../src/domain/use_cases/tasks_use_case.py | 13 +- agentex/tests/integration/test_task_stream.py | 199 +++++++++++++++++- 5 files changed, 276 insertions(+), 22 deletions(-) diff --git a/agentex/src/domain/entities/tasks.py b/agentex/src/domain/entities/tasks.py index 76e93cf9..21949a7d 100644 --- a/agentex/src/domain/entities/tasks.py +++ b/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( ..., diff --git a/agentex/src/domain/services/task_service.py b/agentex/src/domain/services/task_service.py index d5d1291f..26e8437f 100644 --- a/agentex/src/domain/services/task_service.py +++ b/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 diff --git a/agentex/src/domain/use_cases/streams_use_case.py b/agentex/src/domain/use_cases/streams_use_case.py index e25adffa..e8f14da3 100644 --- a/agentex/src/domain/use_cases/streams_use_case.py +++ b/agentex/src/domain/use_cases/streams_use_case.py @@ -5,6 +5,7 @@ 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 @@ -12,8 +13,10 @@ 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)] diff --git a/agentex/src/domain/use_cases/tasks_use_case.py b/agentex/src/domain/use_cases/tasks_use_case.py index 7c114744..8cc619a1 100644 --- a/agentex/src/domain/use_cases/tasks_use_case.py +++ b/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, diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index 7da23edd..918b0916 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,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")