优化状态轮询性能和可观测性

Change-Id: Ia021faa5fdde460c3df814ca9f7a6c7ce4e29066

Author: yinsu.zs

src/code/agent/services/gateway/__init__.py

@@ -1,9 +1,5 @@
 from .gateways import CpuGatewayService, HistoryGatewayService
-from .status import (
-    StatusStorageService,
-    StatusPoller,
-    get_status_storage_service
-)
+from .status import StatusPoller
 from .queue import (
     TaskStatus,
     TaskRequest,
@@ -14,9 +10,7 @@
 __all__ = [
     'CpuGatewayService',
     'HistoryGatewayService',
-    'StatusStorageService',
     'StatusPoller',
-    'get_status_storage_service',
     'TaskStatus',
     'TaskQueue',
     'get_task_queue'

src/code/agent/services/gateway/gateways/history_gateway.py

@@ -286,6 +286,7 @@ def _convert_status_to_history_item(self, status_data, task_id, file_mtime=None)
             data = final_result.get("data", {})
             prompt_id = data.get("prompt_id", task_id)
             results = data.get("results", [])
+            execution_time = data.get("execution_time")  # 获取执行时间
 
             # 构造 ComfyUI 兼容的输出格式
             outputs = {}
@@ -325,15 +326,22 @@ def _convert_status_to_history_item(self, status_data, task_id, file_mtime=None)
                 []   # outputs_to_execute - 要执行的输出节点
             ]
 
+            # 构造 status 对象，包含执行时间（如果有）
+            status_obj = {
+                "status_str": "success",
+                "completed": True,
+                "messages": []
+            }
+            
+            # 添加执行时间到 status 中（与原生 ComfyUI 兼容）
+            if execution_time is not None:
+                status_obj["execution_time"] = execution_time
+            
             return {
                 prompt_id: {
                     "prompt": prompt_structure,
                     "outputs": outputs,
-                    "status": {
-                        "status_str": "success",
-                        "completed": True,
-                        "messages": []
-                    }
+                    "status": status_obj
                 }
             }
 

src/code/agent/services/gateway/queue/task_queue.py

@@ -217,18 +217,15 @@ def _start_status_polling(self, task_id: str):
         为任务启动状态轮询
         """
         try:
-            from services.gateway.status import StatusPoller, get_status_storage_service
+            from services.gateway.status import StatusPoller
 
             with self._poller_lock:
                 if task_id in self._status_pollers:
                     log("DEBUG", f"Task {task_id} already has status poller")
                     return
 
-                storage_service = get_status_storage_service()
-                
                 poller = StatusPoller(
                     task_id=task_id,
-                    storage_service=storage_service,
                     poll_interval=0.5
                 )
 

src/code/agent/services/gateway/status/__init__.py

@@ -1,9 +1,6 @@
-from .storage import StatusStorageService, get_status_storage_service
 from .poller import StatusPoller
 
 __all__ = [
-    'StatusStorageService',
     'StatusPoller',
-    'get_status_storage_service'
 ]
 

src/code/agent/services/gateway/status/poller.py

@@ -7,7 +7,6 @@
 import threading
 from typing import Optional, Callable, Any
 
-from .storage import StatusStorageService
 from utils.logger import log
 
 
@@ -17,21 +16,20 @@ class StatusPoller:
     用于定期从存储中轮询工作流执行状态
     """
 
-    def __init__(self, task_id: str, storage_service: StatusStorageService, poll_interval: float = 1.0):
+    def __init__(self, task_id: str, poll_interval: float = 1.0):
         """
         初始化状态轮询器
         
         Args:
             task_id: 要轮询的任务ID
-            storage_service: 状态存储服务实例
             poll_interval: 轮询间隔，单位秒，默认1.0秒
         """
         self.task_id = task_id
-        self.storage_service = storage_service
         self.poll_interval = poll_interval
         self.is_running = False
         self.thread: Optional[threading.Thread] = None
         self.on_status_update: Optional[Callable[[str, Any], None]] = None
+        self.last_status_count = 0  # 记录已处理的状态数量
 
     def set_status_callback(self, callback: Callable[[str, Any], None]):
         """
@@ -64,7 +62,7 @@ def stop(self):
         log("INFO", f"Stopped polling for task {self.task_id}")
 
     def _poll_loop(self):
-        """轮询循环逻辑（使用增量文件读取）"""
+        """轮询循环逻辑（通过serverless_api接口查询状态）"""
         # 是否已找到任务完成状态
         task_completed = False
 
@@ -73,13 +71,16 @@ def _poll_loop(self):
                 # 时间监控：轮询周期开始
                 poll_start = time.time()
 
-                # 使用增量文件读取，仅获取新增状态
+                # 通过serverless_api接口读取状态
                 read_start = time.time()
-                new_statuses = self.storage_service.get_status_incremental(self.task_id)
+                all_statuses = self._get_status_from_api()
                 read_cost = (time.time() - read_start) * 1000
 
+                # 只处理新增的状态
+                new_statuses = all_statuses[self.last_status_count:]
+                
                 if new_statuses:
-                    log("INFO", f"[Perf][{self.task_id}] File read cost: {read_cost:.1f}ms, got {len(new_statuses)} statuses")
+                    log("INFO", f"[Perf][{self.task_id}] File read cost: {read_cost:.1f}ms, got {len(new_statuses)} new statuses (total: {len(all_statuses)})")
 
                     # 调用回调函数处理新状态
                     callback_start = time.time()
@@ -90,6 +91,9 @@ def _poll_loop(self):
                             # 默认输出到控制台
                             log("DEBUG", f"Task {self.task_id} status update: {json.dumps(status, ensure_ascii=False)}")
 
+                        # 更新已处理计数
+                        self.last_status_count += 1
+                        
                         # 检查是否为任务完成状态
                         if self._is_status_completed(status):
                             task_completed = True
@@ -110,11 +114,23 @@ def _poll_loop(self):
 
         self.is_running = False
 
-        # 清理状态读取缓存
+        # 不再需要清理缓存（因为不使用增量读取缓存）
+        log("DEBUG", f"Polling stopped for task {self.task_id}, processed {self.last_status_count} statuses")
+    
+    def _get_status_from_api(self):
+        """
+        通过调用serverless_api接口获取状态
+        
+        Returns:
+            list: 状态列表
+        """
         try:
-            self.storage_service.cleanup_cache(self.task_id)
+            from services.serverlessapi.serverless_api_service import ServerlessApiService
+            service = ServerlessApiService()
+            return service.get_status_from_store(self.task_id)
         except Exception as e:
-            log("WARNING", f"Error cleaning up cache for task {self.task_id}: {e}")
+            log("ERROR", f"Failed to get status from serverless_api for task {self.task_id}: {e}")
+            return []
 
     def _is_status_completed(self, status: dict) -> bool:
         """