src/core/errors.py

"""
统一错误处理模块

提供层次化的错误类体系，兼容 Gemini API 格式。
支持基于 gRPC 状态码和状态字符串的错误解析。
"""

import json
from typing import Any
from enum import Enum

class ErrorStatus(str, Enum):
    """Vcore AI API 错误状态码 (基于 gRPC 标准)"""
    OK = "OK"                                   # 0
    CANCELLED = "CANCELLED"                     # 1
    UNKNOWN = "UNKNOWN"                         # 2
    INVALID_ARGUMENT = "INVALID_ARGUMENT"       # 3 (400)
    DEADLINE_EXCEEDED = "DEADLINE_EXCEEDED"     # 4 (504)
    NOT_FOUND = "NOT_FOUND"                     # 5 (404)
    ALREADY_EXISTS = "ALREADY_EXISTS"           # 6 (409)
    PERMISSION_DENIED = "PERMISSION_DENIED"     # 7 (403)
    RESOURCE_EXHAUSTED = "RESOURCE_EXHAUSTED"   # 8 (429)
    FAILED_PRECONDITION = "FAILED_PRECONDITION" # 9 (400)
    ABORTED = "ABORTED"                         # 10 (409)
    OUT_OF_RANGE = "OUT_OF_RANGE"               # 11 (400)
    UNIMPLEMENTED = "UNIMPLEMENTED"             # 12 (501)
    INTERNAL = "INTERNAL"                       # 13 (500)
    UNAVAILABLE = "UNAVAILABLE"                 # 14 (503)
    DATA_LOSS = "DATA_LOSS"                     # 15 (500)
    UNAUTHENTICATED = "UNAUTHENTICATED"         # 16 (401)

# gRPC 状态到 HTTP 状态码的映射
GRPC_TO_HTTP: dict[ErrorStatus, int] = {
    ErrorStatus.OK: 200,
    ErrorStatus.CANCELLED: 499,
    ErrorStatus.UNKNOWN: 500,
    ErrorStatus.INVALID_ARGUMENT: 400,
    ErrorStatus.DEADLINE_EXCEEDED: 504,
    ErrorStatus.NOT_FOUND: 404,
    ErrorStatus.ALREADY_EXISTS: 409,
    ErrorStatus.PERMISSION_DENIED: 403,
    ErrorStatus.RESOURCE_EXHAUSTED: 429,
    ErrorStatus.FAILED_PRECONDITION: 400,
    ErrorStatus.ABORTED: 409,
    ErrorStatus.OUT_OF_RANGE: 400,
    ErrorStatus.UNIMPLEMENTED: 501,
    ErrorStatus.INTERNAL: 500,
    ErrorStatus.UNAVAILABLE: 503,
    ErrorStatus.DATA_LOSS: 500,
    ErrorStatus.UNAUTHENTICATED: 401,
}

GRPC_CODE_TO_STATUS: dict[int, ErrorStatus] = {
    0: ErrorStatus.OK,
    1: ErrorStatus.CANCELLED,
    2: ErrorStatus.UNKNOWN,
    3: ErrorStatus.INVALID_ARGUMENT,
    4: ErrorStatus.DEADLINE_EXCEEDED,
    5: ErrorStatus.NOT_FOUND,
    6: ErrorStatus.ALREADY_EXISTS,
    7: ErrorStatus.PERMISSION_DENIED,
    8: ErrorStatus.RESOURCE_EXHAUSTED,
    9: ErrorStatus.FAILED_PRECONDITION,
    10: ErrorStatus.ABORTED,
    11: ErrorStatus.OUT_OF_RANGE,
    12: ErrorStatus.UNIMPLEMENTED,
    13: ErrorStatus.INTERNAL,
    14: ErrorStatus.UNAVAILABLE,
    15: ErrorStatus.DATA_LOSS,
    16: ErrorStatus.UNAUTHENTICATED,
}


def _coerce_error_status(status: str | ErrorStatus | None) -> ErrorStatus:
    if isinstance(status, ErrorStatus):
        return status
    if isinstance(status, str):
        try:
            return ErrorStatus(status)
        except ValueError:
            return ErrorStatus.UNKNOWN
    return ErrorStatus.UNKNOWN

class VcoreError(Exception):
    """Vcore AI 代理错误基类"""
    
    def __init__(
        self,
        message: str,
        code: int | None = None,
        status: str | ErrorStatus | None = None,
        details: dict[str, Any] | None = None,
        upstream_response: str | None = None
    ):
        self.message = message
        
        # 规范化 status
        status_enum = _coerce_error_status(status)
        self.status = status_enum.value
            
        # 规范化 code (HTTP 状态码)
        if code is not None:
            try:
                norm_code = int(code)
            except (TypeError, ValueError):
                norm_code = GRPC_TO_HTTP.get(status_enum, 500)

            # Google/GraphQL 错误里 extensions.status.code 常是 gRPC 数字码（1~16），
            # 不能直接当作 HTTP 状态码发送给 ASGI/uvicorn。
            if norm_code in GRPC_CODE_TO_STATUS and norm_code < 100:
                grpc_status = status_enum if status_enum != ErrorStatus.UNKNOWN else GRPC_CODE_TO_STATUS[norm_code]
                self.status = grpc_status.value
                self.code = GRPC_TO_HTTP.get(grpc_status, 500)
            elif 100 <= norm_code <= 599:
                self.code = norm_code
            else:
                self.code = GRPC_TO_HTTP.get(status_enum, 500)
        else:
            self.code = GRPC_TO_HTTP.get(status_enum, 500)
            
        self.details = details or {}
        self.upstream_response = upstream_response
        super().__init__(message)
    
    def to_Dict(self) -> dict[str, Any]:
        """转换为 Gemini API 兼容的错误响应格式"""
        error_dict: dict[str, Any] = {
            "error": {
                "code": self.code,
                "message": self.message,
                "status": self.status
            }
        }
        if self.details:
            error_dict["error"]["details"] = self.details
        return error_dict
    
    def to_json(self) -> str:
        return json.dumps(self.to_Dict(), ensure_ascii=False)
    
    def to_sse(self) -> bytes:
        return f"data: {self.to_json()}\n\n".encode('utf-8')
    
    @property
    def is_retryable(self) -> bool:
        """判断此错误是否可重试"""
        # 408, 429, 5xx 通常可重试
        if self.code in {408, 429, 500, 502, 503, 504}:
            return True
        # 认证错误在我们的场景中（Token过期）也是可重试的
        if isinstance(self, AuthenticationError):
            return True
        return False

class ClientError(VcoreError):
    """客户端错误 (4xx)"""
    pass

class ServerError(VcoreError):
    """服务端错误 (5xx)"""
    pass

class AuthenticationError(ClientError):
    """认证错误 (401/403)"""
    def __init__(self, message: str = "Authentication failed", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 401, ErrorStatus.UNAUTHENTICATED, details, upstream_response)

class PermissionDeniedError(ClientError):
    """权限拒绝错误 (403)"""
    def __init__(self, message: str = "Permission denied", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 403, ErrorStatus.PERMISSION_DENIED, details, upstream_response)

class InvalidArgumentError(ClientError):
    """参数错误 (400)"""
    def __init__(self, message: str = "Invalid argument", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 400, ErrorStatus.INVALID_ARGUMENT, details, upstream_response)

class NotFoundError(ClientError):
    """资源不存在错误 (404)"""
    def __init__(self, message: str = "Resource not found", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 404, ErrorStatus.NOT_FOUND, details, upstream_response)

class RateLimitError(ClientError):
    """速率限制/资源耗尽错误 (429)"""
    def __init__(self, message: str = "Resource exhausted", details: dict[str, Any] | None = None, retry_after: int | None = None, upstream_response: str | None = None):
        super().__init__(message, 429, ErrorStatus.RESOURCE_EXHAUSTED, details, upstream_response)
        self.retry_after = retry_after

class InternalError(ServerError):
    """内部服务器错误 (500)"""
    def __init__(self, message: str = "Internal server error", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 500, ErrorStatus.INTERNAL, details, upstream_response)

class EmptyResponseError(ServerError):
    """上游返回空响应"""
    def __init__(self, message: str = "Upstream returned empty response", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 502, ErrorStatus.INTERNAL, details, upstream_response)

class UpstreamResponseIncompleteError(ServerError):
    """上游在首包前未返回有效响应结构。"""
    def __init__(self, message: str = "Upstream response structure is incomplete", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 502, ErrorStatus.INTERNAL, details, upstream_response)

class RequestPoolTimeoutError(ServerError):
    """请求池在配置阈值内未能获取上游响应首包并选出 winner。"""
    def __init__(self, message: str = "Request pool timed out before receiving upstream first response", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 504, ErrorStatus.DEADLINE_EXCEEDED, details, upstream_response)

class UpstreamResponseTimeoutError(ServerError):
    """winner 首包后上游长时间没有继续返回原始流块。"""
    def __init__(self, message: str = "Upstream response timed out after winner first response", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 504, ErrorStatus.DEADLINE_EXCEEDED, details, upstream_response)

class UpstreamError(ServerError):
    """上游 API 错误（通用）"""
    def __init__(self, message: str, code: int = 502, status: str | None = None, details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, code, status or ErrorStatus.INTERNAL.value, details, upstream_response)

class UnavailableError(ServerError):
    """服务不可用错误 (503)"""
    def __init__(self, message: str = "Service unavailable", details: dict[str, Any] | None = None, upstream_response: str | None = None):
        super().__init__(message, 503, ErrorStatus.UNAVAILABLE, details, upstream_response)

def raise_for_status(
    code: int | str,
    status: str | None = None,
    message: str = "Unknown error",
    details: dict[str, Any] | None = None,
    upstream_response: str | None = None
) -> VcoreError:
    """
    根据 HTTP 状态码或 gRPC 状态字符串创建对应的错误实例
    """
    # 统一转换 code 为 int，如果失败（如传入了字符串状态）则保持
    try:
        norm_code = int(code)
    except (ValueError, TypeError):
        norm_code = code

    status_enum = _coerce_error_status(status)
    grpc_status = GRPC_CODE_TO_STATUS.get(norm_code) if isinstance(norm_code, int) else None
    effective_status = status_enum if status_enum != ErrorStatus.UNKNOWN else grpc_status

    # 优先根据 gRPC 状态码或状态字符串判断
    # code 为 8 或 429 时代表 RESOURCE_EXHAUSTED
    if effective_status == ErrorStatus.RESOURCE_EXHAUSTED or norm_code == 429:
        return RateLimitError(message, details, upstream_response=upstream_response)
    if effective_status == ErrorStatus.UNAUTHENTICATED or norm_code == 401:
        return AuthenticationError(message, details, upstream_response=upstream_response)
    if effective_status == ErrorStatus.PERMISSION_DENIED or norm_code == 403:
        return PermissionDeniedError(message, details, upstream_response=upstream_response)
    if effective_status == ErrorStatus.INVALID_ARGUMENT or norm_code == 400:
        return InvalidArgumentError(message, details, upstream_response=upstream_response)
    if effective_status == ErrorStatus.NOT_FOUND or norm_code == 404:
        return NotFoundError(message, details, upstream_response=upstream_response)
    if effective_status == ErrorStatus.UNAVAILABLE or norm_code == 503:
        return UnavailableError(message, details, upstream_response=upstream_response)

    if effective_status and effective_status != ErrorStatus.OK:
        http_code = GRPC_TO_HTTP.get(effective_status, 500)
        if 400 <= http_code < 500:
            return ClientError(message, http_code, effective_status, details, upstream_response)
        return ServerError(message, http_code, effective_status, details, upstream_response)

    # 降级到通用的 HTTP 范围判断
    if isinstance(norm_code, int):
        if 400 <= norm_code < 500:
            return ClientError(message, norm_code, status, details, upstream_response)
        return ServerError(message, norm_code, status, details, upstream_response)
    
    return VcoreError(message, status=status, details=details, upstream_response=upstream_response)

def parse_error_response(response_data: str | dict[str, Any] | list[Any]) -> VcoreError | None:
    """
    从上游响应中解析错误 (支持 gRPC 风格的 JSON 响应)
    """
    if isinstance(response_data, str):
        try:
            response_data = json.loads(response_data)
        except json.JSONDecodeError:
            return None

    # 处理数组格式 (GraphQL 风格)
    if isinstance(response_data, list):
        for item in response_data:
            err = parse_error_response(item)
            if err: return err
        return None

    if not isinstance(response_data, dict):
        return None

    # 1. 检查嵌套的 error 字段 (标准 Google API)
    if 'error' in response_data:
        err_obj = response_data['error']
        if isinstance(err_obj, dict):
            return raise_for_status(
                code=err_obj.get('code', 500),
                status=err_obj.get('status'),
                message=err_obj.get('message', 'Unknown error'),
                details=err_obj.get('details'),
                upstream_response=json.dumps(response_data)
            )

    # 2. 检查 GraphQL 风格的 errors 数组
    if 'errors' in response_data:
        errors = response_data['errors']
        if isinstance(errors, list) and len(errors) > 0:
            first_err = errors[0]
            if isinstance(first_err, dict):
                # 优先从 extensions.status 中获取 code 和 message
                ext_status = first_err.get('extensions', {}).get('status', {})
                code = ext_status.get('code') or first_err.get('code', 500)
                status = ext_status.get('status') or first_err.get('status')
                message = ext_status.get('message') or first_err.get('message', 'Unknown error')
                
                return raise_for_status(
                    code=code,
                    status=status,
                    message=message,
                    details=first_err.get('details'),
                    upstream_response=json.dumps(response_data)
                )

    # 3. 检查扁平格式
    if 'code' in response_data or 'status' in response_data or 'message' in response_data:
        return raise_for_status(
            code=response_data.get('code', 500),
            status=response_data.get('status'),
            message=response_data.get('message', 'Unknown error'),
            details=response_data.get('details'),
            upstream_response=json.dumps(response_data)
        )

    return None