Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions agentex/src/domain/entities/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
...,
Expand Down
10 changes: 8 additions & 2 deletions agentex/src/domain/services/task_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
69 changes: 57 additions & 12 deletions agentex/src/domain/use_cases/streams_use_case.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Comment thread
greptile-apps[bot] marked this conversation as resolved.
# 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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)]
13 changes: 8 additions & 5 deletions agentex/src/domain/use_cases/tasks_use_case.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading