Merge remote-tracking branch 'origin/main' into cynthiawang/agx1-618-2a-task-stream-lifecycle-metrics

# Conflicts:
#	agentex/src/domain/use_cases/streams_use_case.py

Author: cyntwang99

agentex-ui/app/api/agentex/[...path]/route.ts

@@ -34,14 +34,26 @@ async function proxy(
 
   const method = req.method.toUpperCase();
   const hasBody = method !== 'GET' && method !== 'HEAD';
-  const upstream = await fetch(target, {
-    method,
-    headers,
-    body: hasBody ? req.body : undefined,
-    redirect: 'manual',
-    // @ts-expect-error `duplex` is required to stream a request body (undici)
-    duplex: 'half',
-  });
+  let upstream: Response;
+  try {
+    upstream = await fetch(target, {
+      method,
+      headers,
+      body: hasBody ? req.body : undefined,
+      redirect: 'manual',
+      // A client disconnect must tear down the upstream request; undici otherwise holds it
+      // open and the upstream handler keeps running.
+      signal: req.signal,
+      // @ts-expect-error `duplex` is required to stream a request body (undici)
+      duplex: 'half',
+    });
+  } catch (error) {
+    // The disconnect above rejects the fetch with the signal's abort reason itself, so identity —
+    // not `error.name` — separates it from a transport failure that coincides with a disconnect.
+    // Next's reason is a `ResponseAborted`, so an `AbortError` check would never fire.
+    if (error === req.signal.reason) return new Response(null, { status: 499 });
+    throw error;
+  }
 
   // Pass the upstream body through unbuffered so SSE / streaming responses work.
   const resHeaders = new Headers(upstream.headers);

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_metrics import (
@@ -128,10 +131,23 @@ async def stream_task_events(
             # 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.
+            # the beginning. Reading the cursor before the status check also
+            # catches a task that goes terminal while we connect.
             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
             now = asyncio.get_running_loop().time()
             # last_message_time drives the keepalive ping and is reset by pings;
             # last_event_time tracks only real data pushes (a ping is not a data
@@ -154,9 +170,33 @@ 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
             # 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
@@ -173,29 +213,38 @@ async def stream_task_events(
                         last_message_time = now
                         last_event_time = now
                         stalled = False
+                        # 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()
+                        now = asyncio.get_running_loop().time()
                         # No data event pushed for the stall window: count the
                         # onset once. Keepalive pings deliberately do not reset
                         # last_event_time, so a persistently quiet stream is
                         # visible even though the connection stays alive.
-                        if (
-                            current_time - last_event_time >= stall_threshold
-                            and not stalled
-                        ):
+                        if now - last_event_time >= stall_threshold and not stalled:
                             stalled = True
                             record_stream_stall()
-                        if current_time - last_message_time >= ping_interval:
+                        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
@@ -238,6 +287,7 @@ 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")
             # Only close what we actually opened, so the active gauge is never
             # decremented for a stream that never incremented it.
@@ -246,7 +296,6 @@ async def stream_task_events(
                     outcome,
                     asyncio.get_running_loop().time() - stream_start_time,
                 )
-            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,