review

Change-Id: Ia767a1d14c78bd1aa162e7a59363fb585a4b3f0c

Author: yinsu.zs

src/code/agent/exceptions/exceptions.py

@@ -10,3 +10,45 @@ class StateTransitionError(CustomError):
     def __init__(self, from_state, to_state):
         msg = f"Illegal state transition: {from_state} -> {to_state}"
         super().__init__(message=msg, code=409)
+
+
+# ========== 任务管理相关异常 ==========
+
+class TaskException(CustomError):
+    """任务管理异常基类"""
+    def __init__(self, message: str, error_code: str = "task_error", code: int = 500):
+        self.error_code = error_code
+        super().__init__(code=code, message=message)
+    
+    def to_dict(self):
+        return {
+            "type": "error",
+            "error_code": self.error_code,
+            "error_message": self.message
+        }
+
+
+class ConfigurationException(TaskException):
+    """配置错误"""
+    def __init__(self, message: str):
+        super().__init__(message=message, error_code="configuration_error", code=500)
+
+
+class InvalidRequestException(TaskException):
+    """无效请求"""
+    def __init__(self, message: str, error_code: str = "invalid_request_error"):
+        super().__init__(message=message, error_code=error_code, code=400)
+
+
+class GPUException(TaskException):
+    """GPU相关异常"""
+    def __init__(self, message: str, status_code: int = 500, 
+                 error_code: str = "gpu_error", original_response: dict = None):
+        super().__init__(message=message, error_code=error_code, code=status_code)
+        self.original_response = original_response
+    
+    def to_dict(self):
+        if self.original_response and isinstance(self.original_response, dict):
+            if "type" in self.original_response or "error" in self.original_response:
+                return self.original_response
+        return super().to_dict()

src/code/agent/services/gateway/handlers/prompt_handler.py

@@ -2,9 +2,11 @@
 Prompt Handler
 处理 /prompt 请求逻辑
 """
+import traceback
 from flask import request, jsonify
 
 from utils.logger import log
+from exceptions.exceptions import TaskException
 
 
 class PromptHandler:
@@ -42,26 +44,35 @@ def handle_post_request(self):
                 }
             }), 400
 
-        # 转发给GPU
-        task_id, result = self.task_manager.forward_to_gpu_async(
-            request_body=request_data,
-            client_id=client_id
-        )
-        
-        # 处理结果
-        if task_id:
+        try:
+            # 转发给GPU
+            task_id, result = self.task_manager.forward_to_gpu_async(
+                request_body=request_data,
+                client_id=client_id
+            )
+            
             # 成功：返回ComfyUI格式
             return jsonify({
                 "prompt_id": task_id,
                 "number": 1,
                 "node_errors": {}
             })
-        else:
-            # 失败：result 是 (status_code, error_type, error_message)
-            status_code, error_type, error_message = result
+            
+        except TaskException as e:
+            log("ERROR", f"[PromptHandler] Task error: {e.message}")
+            # 返回ComfyUI格式的错误
+            return jsonify({
+                "error": {
+                    "type": e.error_code,
+                    "message": e.message
+                }
+            }), e.code
+            
+        except Exception as e:
+            log("ERROR", f"[PromptHandler] Unexpected error: {str(e)}\n{traceback.format_exc()}")
             return jsonify({
                 "error": {
-                    "type": error_type,
-                    "message": error_message
+                    "type": "internal_error",
+                    "message": str(e)
                 }
-            }), status_code
+            }), 500

src/code/agent/services/gateway/handlers/serverless_handler.py

@@ -2,9 +2,11 @@
 Serverless Handler
 处理 /serverless/run 请求逻辑
 """
+import traceback
 from flask import request, jsonify
 
 from utils.logger import log
+from exceptions.exceptions import TaskException
 
 class ServerlessHandler:
     """处理 /serverless/run 请求"""
@@ -35,36 +37,36 @@ def handle_post_request(self):
         invocation_type = request.headers.get('X-Fc-Invocation-Type', '').lower()
         is_async = (invocation_type == 'async')
 
-        if is_async:
-            # 异步模式：立即返回任务ID
-            log("INFO", f"[ServerlessHandler] Processing async request")
-            task_id, result = self.task_manager.forward_to_gpu_async(
-                request_body=body,
-                client_id=client_id
-            )
-            
-            # 处理结果
-            if task_id:
+        try:
+            if is_async:
+                # 异步模式：立即返回任务ID
+                log("INFO", f"[ServerlessHandler] Processing async request")
+                task_id, result = self.task_manager.forward_to_gpu_async(
+                    request_body=body,
+                    client_id=client_id
+                )
+                
                 # 成功：返回Serverless格式
                 return jsonify({
                     "task_id": task_id,
                     "status": "pending"
                 }), 202
             else:
-                # 失败：result 是 (status_code, error_type, error_message)
-                status_code, error_type, error_message = result
-                return jsonify({
-                    "type": "error",
-                    "error_code": error_type,
-                    "error_message": error_message
-                }), status_code
-        else:
-            # 同步模式：等待GPU处理完成
-            log("INFO", f"[ServerlessHandler] Processing sync request")
-            
-            # 调用同步转发，直接传递完整的请求体
-            response_data, status_code = self.task_manager.forward_to_gpu_sync(
-                request_body=body
-            )
+                # 同步模式：等待GPU处理完成
+                log("INFO", f"[ServerlessHandler] Processing sync request")
+                response_data, status_code = self.task_manager.forward_to_gpu_sync(
+                    request_body=body
+                )
+                return jsonify(response_data), status_code
+                
+        except TaskException as e:
+            log("ERROR", f"[ServerlessHandler] Task error: {e.message}")
+            return jsonify(e.to_dict()), e.code
 
-            return jsonify(response_data), status_code
+        except Exception as e:
+            log("ERROR", f"[ServerlessHandler] Unexpected error: {str(e)}\n{traceback.format_exc()}")
+            return jsonify({
+                "type": "error",
+                "error_code": "internal_error",
+                "error_message": str(e)
+            }), 500

src/code/agent/services/gateway/task/task_manager.py

@@ -13,6 +13,11 @@
 import requests
 from flask import request
 from utils.logger import log
+from exceptions.exceptions import (
+    ConfigurationException,
+    InvalidRequestException,
+    GPUException
+)
 
 from .task import TaskStatus, Task
 from .utils.task_manager_util import TaskStatusBroadcaster
@@ -388,7 +393,7 @@ def _cleanup_old_completed_tasks_if_needed(self) -> None:
 
     def forward_to_gpu_async(self, 
                              request_body: dict, 
-                             client_id: str) -> Tuple[Optional[str], Union[object, Tuple[int, str, str]]]:
+                             client_id: str) -> Tuple[str, object]:
         """
         GPU异步转发逻辑
         
@@ -397,44 +402,48 @@ def forward_to_gpu_async(self,
             client_id: 客户端ID(用于WebSocket广播)
             
         Returns:
-            tuple: (task_id, response_or_error)
-                   - task_id: 成功时返回任务ID，失败时返回None
-                   - response_or_error: 成功时返回response对象，失败时返回(status_code, error_type, error_message)
+            tuple: (task_id, response) - 成功时返回任务ID和响应对象
+            
+        Raises:
+            ConfigurationException: GPU URL未配置
+            InvalidRequestException: 缺少任务ID或队列已满
+            GPUException: GPU转发失败或异步调用失败
         """
         request_id = request.headers.get('x-fc-request-id', 'unknown')
 
         # 检查GPU URL配置
         if not self._gpu_function_url:
             log("ERROR", f"[TaskManager][RequestId={request_id}] GPU_FUNCTION_URL not configured")
-            return None, (500, "configuration_error", "GPU_FUNCTION_URL not configured for CPU mode")
+            raise ConfigurationException("GPU_FUNCTION_URL not configured for CPU mode")
 
-        # 提取任务ID(如果找不到会抛出异常)
+        # 提取任务ID
         try:
             task_id = self._extract_task_id_from_request()
         except ValueError as e:
             log("ERROR", f"[TaskManager][RequestId={request_id}] {str(e)}")
-            return None, (500, "missing_task_id", str(e))
+            raise InvalidRequestException(str(e), error_code="missing_task_id")
 
-        # 将任务添加到管理器(用于跟踪和状态管理)
-        # 保存完整的 request_body 作为 prompt_body
+        # 将任务添加到管理器
         try:
             self.submit_task(
                 prompt_body=request_body,
                 client_id=client_id,
                 task_id=task_id
             )
+        except InvalidRequestException:
+            raise
         except Exception as e:
             error_msg = f"Failed to add task to queue: {e}"
             log("ERROR", f"[TaskManager][TaskId={task_id}][RequestId={request_id}] {error_msg}\n{traceback.format_exc()}")
-            return None, (500, "queue_error", error_msg)
+            raise GPUException(error_msg, error_code="task_submit_error")
 
         # 构造GPU URL和headers
         gpu_url = f"{self._gpu_function_url.rstrip('/')}/api/serverless/run"
 
         forward_headers = {
-            'x-fc-async-task-id': task_id,  # 优先使用这个作为 task_id
-            'x-fc-trace-id': task_id,       # GPU的request-id与task-id一致
-            'x-fc-invocation-type': 'Async' # 异步调用
+            'x-fc-async-task-id': task_id,
+            'x-fc-trace-id': task_id,
+            'x-fc-invocation-type': 'Async'
         }
 
         # 复制其他 headers
@@ -444,7 +453,6 @@ def forward_to_gpu_async(self,
             if k_lower not in excluded_headers_lower:
                 forward_headers[k] = v
 
-        # 直接转发完整的请求体（保持原有结构，包括可能的 extra_data）
         # 转发请求到GPU
         try:
             resp = requests.post(
@@ -469,7 +477,8 @@ def forward_to_gpu_async(self,
             except Exception as status_error:
                 log("WARNING", f"[TaskManager][TaskId={task_id}] Failed to update task status to FAILED: {status_error}")
 
-            return None, (500, "gpu_forward_error", error_msg)
+            raise GPUException(error_msg, error_code="gpu_forward_error")
+            
         # 检查 GPU 响应状态
         if resp.status_code != 202:
             # GPU 拒绝了请求，更新为 FAILED
@@ -481,47 +490,44 @@ def forward_to_gpu_async(self,
                 }, TaskStatus.FAILED)
             except Exception as e:
                 log("WARNING", f"[TaskId={task_id}] Failed to update task status to FAILED: {e}\n{traceback.format_exc()}")
+            
+            raise GPUException(
+                f"Failed to invoke GPU function asynchronously: HTTP {resp.status_code}",
+                status_code=resp.status_code,
+                error_code="async_invocation_error"
+            )
 
-        # 返回response和task_id
-        if resp.status_code == 202:
-            return task_id, resp
-        else:
-            return None, (500, "async_invocation_error", f"Failed to invoke GPU function asynchronously: HTTP {resp.status_code}")
+        # 返回task_id和response
+        return task_id, resp
 
-    def forward_to_gpu_sync(self,
-                           request_body: dict) -> Tuple[dict, int]:
+    def forward_to_gpu_sync(self, request_body: dict) -> Tuple[dict, int]:
         """
         GPU同步转发逻辑（等待GPU处理完成并返回结果）
         
         Args:
             request_body: 完整的请求体
             
         Returns:
-            tuple: (response_dict, status_code)
-                   - response_dict: 响应数据
-                   - status_code: HTTP状态码
+            tuple: (response_dict, status_code) - GPU返回的原始响应
+            
+        Raises:
+            ConfigurationException: GPU URL未配置
+            InvalidRequestException: 缺少任务ID
+            GPUException: GPU转发失败或返回非2xx状态码
         """
         request_id = request.headers.get('x-fc-request-id', 'unknown')
 
         # 检查GPU URL配置
         if not self._gpu_function_url:
             log("ERROR", f"[TaskManager][RequestId={request_id}] GPU_FUNCTION_URL not configured")
-            return {
-                "type": "error",
-                "error_code": "configuration_error",
-                "error_message": "GPU_FUNCTION_URL not configured for CPU mode"
-            }, 500
+            raise ConfigurationException("GPU_FUNCTION_URL not configured for CPU mode")
 
-        # 提取任务ID(如果找不到会抛出异常)
+        # 提取任务ID
         try:
             task_id = self._extract_task_id_from_request()
         except ValueError as e:
             log("ERROR", f"[TaskManager][RequestId={request_id}] {str(e)}")
-            return {
-                "type": "error",
-                "error_code": "missing_task_id",
-                "error_message": str(e)
-            }, 500
+            raise InvalidRequestException(str(e), error_code="missing_task_id")
 
         # 构造GPU URL和headers
         gpu_url = f"{self._gpu_function_url.rstrip('/')}/api/serverless/run"
@@ -546,34 +552,34 @@ def forward_to_gpu_sync(self,
                 json=request_body,
                 headers=forward_headers,
                 params=request.args,
-                timeout=600  # 同步调用使用更长的超时时间（10分钟）
+                timeout=600
             )
 
             if resp.status_code == 200:
                 log("INFO", f"[TaskManager][TaskId={task_id}][RequestId={request_id}] Sync request completed successfully")
                 return resp.json(), 200
             else:
                 error_msg = f"GPU returned HTTP {resp.status_code}"
+                log("ERROR", f"[TaskManager][TaskId={task_id}][RequestId={request_id}] Sync request failed: {error_msg}")
+                
+                # 尝试解析GPU的错误响应
                 try:
                     error_response = resp.json()
-                    log("ERROR", f"[TaskManager][TaskId={task_id}][RequestId={request_id}] Sync request failed: {error_msg}, response: {error_response}")
-                    return error_response, resp.status_code
-                except Exception:
-                    log("ERROR", f"[TaskManager][TaskId={task_id}][RequestId={request_id}] Sync request failed: {error_msg}")
-                    return {
-                        "type": "error",
-                        "error_code": "gpu_error",
-                        "error_message": error_msg
-                    }, resp.status_code
+                    raise GPUException(
+                        error_msg,
+                        status_code=resp.status_code,
+                        original_response=error_response
+                    )
+                except (ValueError, requests.JSONDecodeError):
+                    raise GPUException(error_msg, status_code=resp.status_code)
+                    
+        except GPUException:
+            raise
         except Exception as e:
             # 处理转发GPU失败的情况
             error_msg = f"Failed to send sync request to GPU: {str(e)}"
             log("ERROR", f"[TaskManager][TaskId={task_id}][RequestId={request_id}] {error_msg}\nStacktrace:\n{traceback.format_exc()}")
-            return {
-                "type": "error",
-                "error_code": "gpu_forward_error",
-                "error_message": error_msg
-            }, 500
+            raise GPUException(error_msg, error_code="gpu_forward_error")
 
     @staticmethod
     def _extract_task_id_from_request() -> str: