fix

Author: rayrayraykk

cookbook/en/agent_app.md

@@ -435,10 +435,12 @@ Status query response (completed):
 
 **Important Notes**
 
-1. **In-memory mode**: By default uses in-memory storage; task state is lost on process restart
-2. **Persistence**: For production, configure Celery with Redis for task persistence
-3. **Storage**: Only stores final response; intermediate streaming events are not saved
-4. **Timeout**: Set reasonable timeout based on agent complexity
+1. **Dual mode support**:
+   - **In-memory mode** (default): Task state is lost on restart; suitable for development/testing
+   - **Celery mode**: Configure `broker_url` and `backend_url` to enable; tasks persisted; suitable for production
+2. **Storage**: Only stores final response; intermediate streaming events are not saved
+3. **Timeout**: Set reasonable timeout based on agent complexity
+4. **Worker requirement**: Celery mode requires running workers (use `enable_embedded_worker=True`)
 
 ------
 

cookbook/zh/agent_app.md

@@ -440,10 +440,12 @@ while True:
 
 **注意事项**
 
-1. **内存模式**：默认使用 in-memory 模式，进程重启后任务状态会丢失
-2. **持久化**：生产环境建议配置 Celery（需要 Redis）实现任务持久化
-3. **结果存储**：只存储最终 response，中间流式事件不会被保存
-4. **超时设置**：建议根据 agent 复杂度设置合理的超时时间
+1. **双模式支持**：
+   - **In-memory 模式**（默认）：进程重启后任务状态会丢失，适合开发和测试
+   - **Celery 模式**：配置 `broker_url` 和 `backend_url` 启用，任务持久化，适合生产环境
+2. **结果存储**：只存储最终 response，中间流式事件不会被保存
+3. **超时设置**：建议根据 agent 复杂度设置合理的超时时间
+4. **Worker 需求**：Celery 模式需要启动 worker（可使用 `enable_embedded_worker=True`）
 
 ------
 

pyproject.toml

@@ -62,6 +62,7 @@ dev = [
     "sphinx-autoapi>=3.6.0",
     "pytest-mock>=3.15.1",
     "sphinxcontrib-mermaid>=1.2.3",
+    "aiohttp>=3.9.0",
 ]
 
 ext = [

src/agentscope_runtime/engine/app/agent_app.py

@@ -7,7 +7,9 @@
 import platform
 import shlex
 import subprocess
+import time
 import types
+import uuid
 from contextlib import asynccontextmanager, AsyncExitStack
 from typing import Any, Callable, Dict, List, Optional, Type
 
@@ -173,6 +175,7 @@ def __init__(
         self.enable_stream_task = enable_stream_task
         self.stream_task_queue = stream_task_queue
         self.stream_task_timeout = stream_task_timeout
+        self._stream_query_celery_task: Optional[Callable] = None
 
         self._query_handler: Optional[Callable] = None
         self._init_handler: Optional[Callable] = None
@@ -249,6 +252,7 @@ async def _internal_framework_lifespan(self, app: FastAPI):
         """
         # pylint: disable=too-many-branches
         self._build_runner()
+        cleanup_task = None
         try:
             # aexit any possible running instances before set up
             # runner
@@ -272,9 +276,22 @@ async def _internal_framework_lifespan(self, app: FastAPI):
             if self.enable_embedded_worker and self.celery_app:
                 self.start_embedded_celery_worker()
 
+            if self.enable_stream_task:
+                cleanup_task = asyncio.create_task(
+                    self._task_cleanup_worker(),
+                )
+                logger.info("Started task cleanup worker")
+
             yield
 
         finally:
+            if cleanup_task:
+                cleanup_task.cancel()
+                try:
+                    await cleanup_task
+                except asyncio.CancelledError:
+                    pass
+
             if self.after_finish:
                 try:
                     if asyncio.iscoroutinefunction(self.after_finish):
@@ -406,12 +423,92 @@ async def root():
 
         self._add_process_control_endpoints()
 
-    def _add_stream_query_task_endpoint(self):
+    async def _cleanup_expired_tasks(self):
+        """
+        Remove completed/failed tasks older than TTL.
+
+        Returns:
+            Number of tasks cleaned up
+        """
+        now = time.time()
+        ttl_seconds = 3600  # 1 hour
+
+        expired = []
+        for task_id, info in self.active_tasks.items():
+            status = info.get("status")
+
+            if status in ["completed", "failed"]:
+                finished_at = info.get("completed_at") or info.get(
+                    "failed_at",
+                )
+                if finished_at and (now - finished_at) > ttl_seconds:
+                    expired.append(task_id)
+
+        for task_id in expired:
+            del self.active_tasks[task_id]
+            if hasattr(self, "task_locks") and task_id in self.task_locks:
+                del self.task_locks[task_id]
+
+        if expired:
+            logger.info(
+                f"Cleaned up {len(expired)} expired tasks. "
+                f"Active tasks: {len(self.active_tasks)}",
+            )
+
+        return len(expired)
+
+    async def _task_cleanup_worker(self):
+        """Background worker to cleanup expired tasks periodically."""
+        while True:
+            try:
+                await asyncio.sleep(300)  # Run every 5 minutes
+                await self._cleanup_expired_tasks()
+            except asyncio.CancelledError:
+                logger.info("Task cleanup worker stopped")
+                break
+            except Exception as e:
+                logger.error(f"Task cleanup failed: {e}")
+
+    def _create_stream_query_wrapper(self):
+        """
+        Create a wrapper function for stream_query that collects only
+        the final response.
+
+        This wrapper is used by Celery to execute stream_query as a
+        background task.
+        """
+
+        async def stream_query_wrapper(request: dict):
+            """Wrapper that collects only final response from stream_query"""
+            final_response = None
+
+            async for event in self._runner.stream_query(request):
+                if hasattr(event, "model_dump"):
+                    final_response = event.model_dump()
+                elif hasattr(event, "dict"):
+                    final_response = event.dict()
+                else:
+                    final_response = {"data": str(event)}
+
+            return final_response
+
+        return stream_query_wrapper
+
+    def _add_stream_query_task_endpoint(self) -> None:
         """
         Add background task endpoints for stream_query.
-        Creates POST /process/task and GET /process/task/{task_id}.
 
+        Creates POST /process/task and GET /process/task/{task_id}.
         Design: Only stores the final response, not intermediate events.
+        Supports both Celery and in-memory modes.
+
+        Args:
+            self (AgentApp): The application instance on which to register
+                the task endpoints.
+
+        Returns:
+            None: This method registers routes on the application and does
+                not return a value.
         """
         if not self.enable_stream_task:
             logger.debug("Stream task disabled, skipping task endpoint setup")
@@ -447,22 +544,27 @@ def _add_stream_query_task_endpoint(self):
         @UnifiedRoutingMixin.internal_route
         async def submit_stream_query_task(request: dict):
             """Submit stream_query as background task"""
-            import uuid
-
             task_id = str(uuid.uuid4())
 
             if self.celery_app:
+                if self._stream_query_celery_task is None:
+                    wrapper_func = self._create_stream_query_wrapper()
+                    self._stream_query_celery_task = self.register_celery_task(
+                        wrapper_func,
+                        self.stream_task_queue,
+                    )
+
+                result = self._stream_query_celery_task.delay(request)
+
                 return {
-                    "error": (
-                        "Celery mode not yet implemented for stream tasks"
-                    ),
-                    "suggestion": (
-                        "Use in-memory mode or contribute implementation"
+                    "task_id": result.id,
+                    "status": "submitted",
+                    "queue": self.stream_task_queue,
+                    "message": (
+                        "Stream query task submitted to Celery successfully"
                     ),
                 }
             else:
-                import time
-
                 self.active_tasks[task_id] = {
                     "task_id": task_id,
                     "status": "submitted",

src/agentscope_runtime/engine/deployers/utils/service_utils/routing/task_engine_mixin.py

@@ -19,6 +19,8 @@ def init_task_engine(
         self.celery_app = None
         self.active_tasks: Dict[str, Dict[str, Any]] = {}
         self._registered_queues: set[str] = set()
+        self.task_locks: Dict[str, asyncio.Lock] = {}
+        self._tasks_lock: Optional[asyncio.Lock] = None
 
         if broker_url and backend_url:
             try:
@@ -156,6 +158,24 @@ def _run_celery_task_processor(
 
         self.celery_app.worker_main(cmd)
 
+    async def _get_task_lock(self, task_id: str) -> asyncio.Lock:
+        """
+        Get or create a lock for a specific task.
+
+        Args:
+            task_id: Task identifier
+
+        Returns:
+            asyncio.Lock for the specified task
+        """
+        if self._tasks_lock is None:
+            self._tasks_lock = asyncio.Lock()
+
+        async with self._tasks_lock:
+            if task_id not in self.task_locks:
+                self.task_locks[task_id] = asyncio.Lock()
+            return self.task_locks[task_id]
+
     async def execute_background_task(
         self,
         task_id: str,
@@ -242,54 +262,89 @@ async def execute_stream_query_task(
 
         Returns:
             Final response event as dict
+
+        Raises:
+            TimeoutError: If task exceeds specified timeout
+            RuntimeError: If stream yields no events
         """
         # pylint:disable=unused-argument
+        task_lock = await self._get_task_lock(task_id)
+
         try:
-            self.active_tasks[task_id].update(
-                {
-                    "status": "running",
-                    "started_at": time.time(),
-                },
-            )
+            async with task_lock:
+                self.active_tasks[task_id].update(
+                    {
+                        "status": "running",
+                        "started_at": time.time(),
+                    },
+                )
 
             final_response = None
             start_time = time.time()
+            event_count = 0
+
+            async def stream_with_collection():
+                nonlocal final_response, event_count
+                async for event in stream_func(request):
+                    event_count += 1
+
+                    if hasattr(event, "model_dump"):
+                        final_response = event.model_dump()
+                    elif hasattr(event, "dict"):
+                        final_response = event.dict()
+                    else:
+                        final_response = {"data": str(event)}
+
+            if timeout is not None:
+                await asyncio.wait_for(
+                    stream_with_collection(),
+                    timeout=timeout,
+                )
+            else:
+                await stream_with_collection()
 
-            async for event in stream_func(request):
-                if timeout and (time.time() - start_time) > timeout:
-                    raise TimeoutError(
-                        f"Task {task_id} exceeded timeout of {timeout}s",
-                    )
-
-                if hasattr(event, "model_dump"):
-                    final_response = event.model_dump()
-                elif hasattr(event, "dict"):
-                    final_response = event.dict()
-                else:
-                    final_response = {"data": str(event)}
+            if event_count == 0 or final_response is None:
+                raise RuntimeError(
+                    f"Stream function yielded no events for task {task_id}",
+                )
 
             elapsed = time.time() - start_time
 
-            self.active_tasks[task_id].update(
-                {
-                    "status": "completed",
-                    "result": final_response,
-                    "completed_at": time.time(),
-                    "elapsed_time": elapsed,
-                },
-            )
+            async with task_lock:
+                self.active_tasks[task_id].update(
+                    {
+                        "status": "completed",
+                        "result": final_response,
+                        "completed_at": time.time(),
+                        "elapsed_time": elapsed,
+                    },
+                )
 
             return final_response
 
+        except asyncio.TimeoutError:
+            async with task_lock:
+                self.active_tasks[task_id].update(
+                    {
+                        "status": "failed",
+                        "error": f"Task exceeded timeout of {timeout}s",
+                        "error_type": "TimeoutError",
+                        "failed_at": time.time(),
+                    },
+                )
+            raise
+
         except Exception as e:
-            self.active_tasks[task_id].update(
-                {
-                    "status": "failed",
-                    "error": str(e),
-                    "error_type": type(e).__name__,
-                    "failed_at": time.time(),
-                },
-            )
+            async with task_lock:
+                self.active_tasks[task_id].update(
+                    {
+                        "status": "failed",
+                        "error": str(e),
+                        "error_type": type(e).__name__,
+                        "failed_at": time.time(),
+                    },
+                )
+            raise
 
     def get_task_status(self, task_id: str):
         # pylint:disable=too-many-return-statements