review

Change-Id: I027ff1515728fa6859385edb49cc477cee343c64

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/routes/gateway_routes.py

@@ -5,6 +5,7 @@
 import time
 import traceback
 
+import requests
 from flask import Blueprint, Flask, jsonify, request
 from flask_sock import Sock
 import websocket
@@ -67,6 +68,7 @@ def setup_routes(self):
             self._register_queue_handler()
             self._register_prompt_handler()
             self._register_serverless_run_handler()
+            self._register_serverless_status_handler()
             self._register_history_handler()
             # 通过环境变量控制是否禁用工作流保存
             if constants.DISABLE_FLOW_SAVE:
@@ -213,6 +215,53 @@ def handle_serverless_run():
             """
             return self.serverless_handler.handle_post_request()
 
+    def _register_serverless_status_handler(self):
+        @self.bp.route("/serverless/status", methods=["GET"])
+        @handle_exceptions(error_type="serverless_status_error", log_prefix="ServerlessStatus")
+        def handle_serverless_status():
+            """
+            处理 /api/serverless/status 请求
+            
+            转发到 GPU 函数获取任务状态
+            """
+            task_id = request.args.get("task_id", "")
+            if not task_id:
+                return jsonify({
+                    "type": "error",
+                    "error_code": "invalid_params",
+                    "error_message": "task_id is required"
+                }), 400
+            
+            # 构造 GPU URL
+            gpu_function_url = constants.GPU_FUNCTION_URL
+            if not gpu_function_url:
+                return jsonify({
+                    "type": "error",
+                    "error_code": "configuration_error",
+                    "error_message": "GPU_FUNCTION_URL not configured"
+                }), 500
+            
+            gpu_url = f"{gpu_function_url.rstrip('/')}/api/serverless/status"
+            
+            try:
+                # 转发请求到 GPU
+                resp = requests.get(
+                    gpu_url,
+                    params={"task_id": task_id},
+                    headers=dict(request.headers),
+                    timeout=30
+                )
+                
+                # 返回 GPU 的响应
+                return resp.json(), resp.status_code
+            except Exception as e:
+                log("ERROR", f"[ServerlessStatus] Failed to forward status request: {e}")
+                return jsonify({
+                    "type": "error",
+                    "error_code": "gpu_forward_error",
+                    "error_message": f"Failed to get status from GPU: {str(e)}"
+                }), 500
+    
     def _register_history_handler(self):
         @self.bp.route("/history", methods=["GET"])
         @handle_exceptions(error_type="history_operation_error", log_prefix="History")

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