Spaces:

ZHIWEI666
/

ComfyUI-Ranking-API

Running

Upload 5 files
f1b05a2 verified 3 days ago
24.1 kB
# ⚙️ 后端逻辑/核心服务端.py (Hugging Face Spaces app.py)
from runtime_log_privacy import install_uvicorn_privacy_filters, route_label, RequestTraceMiddleware
install_uvicorn_privacy_filters()

from fastapi import FastAPI, File, UploadFile, Form, Depends, HTTPException, Request
from rate_limit import shared_limit
from response_privacy import PrivateResponseCacheControl
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, JSONResponse, FileResponse 
from pydantic import BaseModel
from huggingface_hub import hf_hub_download, HfApi
import uuid
import urllib.parse
import urllib.request
import urllib.error
import os
import logging
import time
import threading
import 数据库连接 as db
import asyncio
from 安全认证 import require_auth
from profile_io import start_profile_io, close_profile_io
from verification_io import start_verification_io, stop_verification_io, close_verification_io
from auth_io import start_auth_io, stop_auth_io, close_auth_io
from auth_io import start_identity_io, stop_identity_io, close_identity_io
from rate_limit import start_shared_work, stop_shared_work, close_shared_work
from script_execution import start_admin_scripts, stop_admin_scripts, close_admin_scripts
from comment_store import start_comment_reads, stop_comment_reads, close_comment_reads
from periodic_tasks import start_periodic_task, request_periodic_stop, wait_periodic_tasks
from owned_thread import run_owned_thread
from resource_work import start_resource_work, stop_resource_work, close_resource_work
from shutdown_lifecycle import (prepare_application_shutdown, run_application_shutdown,
                                shutdown_cancel_requested, ShutdownFailures)
from startup_warmup import (prepare_startup_warmup, run_startup_warmup, stop_startup_warmup,
                            wait_startup_warmup, startup_warmup_stopped, startup_shutdown_requested)


# ==========================================
# 📝 P2优化：日志配置
# ==========================================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(name)s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("ComfyUI-Ranking")
_UPLOAD_CONCURRENCY = max(1, min(int(os.environ.get("MAX_CONCURRENT_UPLOADS", "4")), 16))
_UPLOAD_SEMAPHORE = threading.BoundedSemaphore(_UPLOAD_CONCURRENCY)


from 云端_定时版本检测引擎 import daily_version_check_task, reset_daily_views_task

# ==========================================
# 👥 用户模块 (拆分为3个子模块)
# ==========================================
# router_users_auth.py    - 🔐 登录/注册/密码重置/验证码
# router_users_profile.py - 👤 获取/更新用户资料
# router_users_social.py  - 🤝 关注/隐私设置
from router_users_auth import router as users_auth_router
from router_users_profile import router as users_profile_router
from router_users_social import router as users_social_router

# ==========================================
# 其他业务模块
# ==========================================
from router_items import (router as items_router, start_resource_reads, stop_resource_reads, close_resource_reads)       # 📦 内容管理(工具/Skill/应用/推荐)
from router_comments import (router as comments_router, start_comment_writes, stop_comment_writes, close_comment_writes) # 💬 评论系统
from router_messages import (router as messages_router, start_message_io, stop_message_io, close_message_io) # ✉️ 私信系统
from router_wallet import (router as wallet_router, start_wallet_sql, stop_wallet_sql, close_wallet_sql)     # 💰 钱包/提现
from proxy_access import start_proxy_access, stop_proxy_access, close_proxy_access
from proxy_download_work import start_proxy_downloads, stop_proxy_downloads, close_proxy_downloads
from media_proxy_service import start_public_media, stop_public_media, close_public_media, public_media_endpoint
from router_proxy import router as proxy_router       # 🔗 代理下载
from router_posts import (router as posts_router, start_posts_io, stop_posts_io, close_posts_io)       # 💬 讨论区
from router_prompts import (router as prompts_router, start_prompts_io, stop_prompts_io, close_prompts_io)   # ✨ 提示词市场
from router_tasks import (router as tasks_router, start_tasks_io, stop_tasks_io, close_tasks_io)       # 📝 任务榜
from router_open_api import start_open_sql, stop_open_sql, close_open_sql
from router_open_api import router as open_api_router  # 🌐 通用开放API（第三方插件）
from router_admin_plugins import router as admin_plugins_router  # 🔒 管理后台 - 第三方插件管理API

from database_sql import init_sql_db
from router_system_config import (router as system_config_router, start_config_sql, stop_config_sql, close_config_sql)

from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException

app = FastAPI(
    title="ComfyUI Ranking API",
    description="ComfyUI 社区排名系统 API",
    version="2.0.0",
    docs_url="/api/docs",
    redoc_url="/api/redoc"
)


# ==========================================
# 🛡️ 稳定性优化：全局异常处理器
# ==========================================
# 作用：捕获所有未处理异常，防止单个请求崩溃整个服务
# 🚀 P3优化：统一错误响应格式

def create_error_response(status_code: int, detail: str, error_type: str = None, error_id: str = None, extra: dict = None):
    """🚀 P3优化：统一错误响应构造器"""
    response = {
        "success": False,
        "status": "error",
        "code": status_code,
        "detail": detail,
    }
    if error_type:
        response["type"] = error_type
    if error_id:
        response["error_id"] = error_id
    if extra:
        response.update(extra)
    return response

@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request, exc):
    """🚀 P3优化：HTTP 异常处理，返回统一格式"""
    logger.warning(f"HTTP {exc.status_code} | {route_label(request)} | detail_type={type(exc.detail).__name__}")
    
    # 根据状态码确定错误类型
    error_type_map = {
        400: "bad_request",
        401: "unauthorized",
        403: "forbidden",
        404: "not_found",
        409: "conflict",
        429: "rate_limited",
        502: "upstream_error",
    }
    error_type = error_type_map.get(exc.status_code, "http_error")
    
    return JSONResponse(
        status_code=exc.status_code,
        content=create_error_response(exc.status_code, str(exc.detail), error_type),
        headers={"X-Ranking-Error-Scope": "endpoint"} if exc.status_code == 503 else None,
    )

from request_body_guard import RequestBodyGuard
app.add_middleware(RequestBodyGuard, error_response=create_error_response)

from upload_body_guard import UploadBodyGuard
app.add_middleware(UploadBodyGuard, capacity=_UPLOAD_CONCURRENCY, error_response=create_error_response)

from public_list_etag import PublicListETag
app.add_middleware(PublicListETag)

app.add_middleware(PrivateResponseCacheControl)


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request, exc):
    """🚀 P3优化：请求参数校验失败处理"""
    logger.warning("Validation Error | %s | count=%d", route_label(request), len(exc.errors()))
    
    # 提取简洁的错误信息
    errors = []
    for err in exc.errors():
        field = ".".join(str(loc) for loc in err["loc"][1:])  # 跳过 'body'
        errors.append({"field": field, "message": "该字段格式不正确", "type": err["type"]})
    
    return JSONResponse(
        status_code=422,
        content=create_error_response(
            422, 
            "请求参数格式错误", 
            "validation_error",
            extra={"errors": errors}
        )
    )

@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    """🚀 P3优化：全局异常处理，捕获所有未预期异常"""
    error_id = getattr(request.state, 'ranking_request_id', None) or uuid.uuid4().hex
    logger.error(f"Unhandled Exception [{error_id}] | {route_label(request)} | error_type={type(exc).__name__}")
    return JSONResponse(
        status_code=500,
        content=create_error_response(
            500,
            "服务器内部错误，请稍后重试",
            "internal_error",
            error_id=error_id
        ),
        headers={"X-Ranking-Request-Id": error_id}
    )


# ==========================================
# ❤️ 健康检查接口 (增强版)
# ==========================================
@app.get("/")
def health_check():
    return {"status": "ok", "message": "ComfyUI Ranking API is running perfectly!"}

@app.get("/health")
def detailed_health_check():
    """详细健康检查，含依赖状态"""
    health_status = {
        "status": "ok",
        "components": {}
    }
    
    # 检查 SQL 数据库
    try:
        from database_sql import check_db_connection
        if check_db_connection():
            health_status["components"]["sql_database"] = "ok"
        else:
            health_status["components"]["sql_database"] = "error: connection failed"
            health_status["status"] = "degraded"
    except Exception as e:
        logger.warning(f"健康检查 SQL 数据库异常: error_type={type(e).__name__}")
        health_status["components"]["sql_database"] = "error: connection failed"
        health_status["status"] = "degraded"
    
    # 检查 JSON 文件访问
    try:
        import 数据库连接 as json_db
        json_db.load_data("users.json", default_data={})
        health_status["components"]["json_storage"] = "ok"
    except Exception as e:
        logger.warning(f"健康检查 JSON 存储异常: error_type={type(e).__name__}")
        health_status["components"]["json_storage"] = "error: storage unavailable"
        health_status["status"] = "degraded"
    
    return health_status

def _warmup_dependencies():
    if startup_warmup_stopped(app.state):
        return
    # ========== 预热检查：SQL 数据库 ==========
    try:
        init_sql_db()
        logger.info("✅ SQL 数据库初始化完成")
    except Exception as e:
        logger.error(f"❌ SQL 数据库初始化失败: error_type={type(e).__name__}")
        # 不抛出异常，允许降级运行

    if startup_warmup_stopped(app.state):
        return
    # ========== 加载第三方插件注册表（依赖 SQL 数据库） ==========
    try:
        from router_open_api import load_plugins_from_db
        load_plugins_from_db()
        logger.info("✅ 第三方插件注册表已加载")
    except Exception as e:
        logger.warning(f"⚠️ 第三方插件注册表加载失败: error_type={type(e).__name__}")

    if startup_warmup_stopped(app.state):
        return
    # ========== 预热检查：JSON 数据库 ==========
    try:
        db.load_data("users.json", default_data={})
        if startup_warmup_stopped(app.state):
            return
        db.load_data("items.json", default_data=[])
        logger.info("✅ JSON 数据库访问正常")
    except Exception as e:
        logger.warning(f"⚠️ JSON 数据库访问异常: error_type={type(e).__name__}")


# ==========================================
# 🚀 应用启动事件
# ==========================================
# 作用：初始化数据库 + 启动定时版本检测后台任务 + 预热检查
@app.on_event("startup")
async def on_startup():
    logger.info("🚀 ComfyUI-Ranking API 启动中...")
    if getattr(app.state, "_ranking_starting", False) or getattr(app.state, "_ranking_closers", 0):
        raise RuntimeError("application lifecycle is already in progress")
    if getattr(app.state, "_ranking_started", False):
        return
    app.state._ranking_starting = True
    try:
        prepare_startup_warmup(app.state)
        prepare_application_shutdown(app.state)
        start_profile_io()
        try:
            start_verification_io()
            start_auth_io()
            start_identity_io()
            start_shared_work()
            start_open_sql()
            start_config_sql()
            start_wallet_sql()
            start_resource_work()
            start_resource_reads()
            start_comment_writes()
            start_message_io()
            start_posts_io()
            start_prompts_io()
            start_tasks_io()
            start_proxy_access()
            start_proxy_downloads()
            start_public_media()
            start_comment_reads()
            start_admin_scripts()

            await run_startup_warmup(app.state, _warmup_dependencies)

            # SQL资金已冻结但JSON发布失败的任务，重启后继续投影；持有句柄供停机收尾。
            from task_publication import publication_worker
            from database_sql import engine
            from db_utils import sort_cache
            start_periodic_task(app.state, "task_publication",
                lambda stop: publication_worker(engine, db, sort_cache.invalidate, stop))

            from resource_publication import publication_worker as resource_publication_worker
            start_periodic_task(app.state, "resource_publication",
                lambda stop: resource_publication_worker(engine, db, sort_cache.invalidate, stop))

            from resource_deletion import resource_deletion_worker
            start_periodic_task(app.state, "resource_deletion",
                lambda stop: resource_deletion_worker(db, stop, sort_cache.invalidate))

            # ========== 启动后台任务 ==========
            start_periodic_task(app.state, "daily_version", daily_version_check_task)
            logger.info("✅ 定时版本检测任务已挂载")

            start_periodic_task(app.state, "daily_views", reset_daily_views_task)
            logger.info("✅ daily_views 每日重置任务已挂载")

            # ========== 启动 HF 批量同步定时器 ==========
            db.start_batch_sync()
            logger.info("✅ HF 同步调度初始化已完成")

            app.state._ranking_started = True
            logger.info("🎉 ComfyUI-Ranking API 启动完成!")
        except BaseException:
            # 启动可能只登记了部分句柄；正常关闭按已登记对象收尾。
            if not startup_shutdown_requested(app.state):
                try:
                    await on_shutdown()
                except BaseException as error:
                    logger.error("启动收尾失败: error_type=%s", type(error).__name__)
            raise
    finally:
        app.state._ranking_starting = False


@app.on_event("shutdown")
async def on_shutdown():
    return await run_application_shutdown(app.state, _shutdown_resources)


async def _shutdown_resources():
    """逐项收尾，错误在其它组件得到关闭机会后传播。"""
    app.state._ranking_started = False
    app.state._ranking_shutdown_incomplete = True
    failures = ShutdownFailures(logger)
    failures.call('log_start', logger.info, "🛑 ComfyUI-Ranking API 正在关闭...")
    failures.call('stop_warmup', stop_startup_warmup, app.state)
    for stage, stop in (
        ('stop_wallet_sql', stop_wallet_sql),
        ('stop_proxy_downloads', stop_proxy_downloads),
        ('stop_public_media', stop_public_media),
        ('stop_resource_reads', stop_resource_reads), ('stop_comment_writes', stop_comment_writes),
        ('stop_posts', stop_posts_io), ('stop_prompts', stop_prompts_io), ('stop_tasks', stop_tasks_io), ('stop_proxy_access', stop_proxy_access),
        ('stop_messages', stop_message_io), ('stop_resources', stop_resource_work), ('stop_open_sql', stop_open_sql), ('stop_config_sql', stop_config_sql), ('stop_shared', stop_shared_work),
        ('stop_identity', stop_identity_io), ('stop_auth', stop_auth_io),
        ('stop_verification', stop_verification_io), ('stop_scripts', stop_admin_scripts),
        ('stop_comments', stop_comment_reads),
    ):
        failures.call(stage, stop)
    workers = []
    for name in ('daily_version', 'daily_views', 'task_publication', 'resource_publication', 'resource_deletion'):
        worker = failures.call('capture_' + name, getattr, app.state, name + '_worker', None)
        if worker is not None:
            workers.append(worker)
        failures.call('stop_' + name, request_periodic_stop, app.state, (name,))
    for stage, close, args in (
        ('wait_warmup', wait_startup_warmup, (app.state,)),
        ('close_proxy_downloads', close_proxy_downloads, ()),
        ('close_public_media', close_public_media, ()),
        ('close_wallet_sql', close_wallet_sql, ()),
        ('close_resource_reads', close_resource_reads, ()), ('close_comment_writes', close_comment_writes, ()),
        ('close_posts', close_posts_io, ()), ('close_prompts', close_prompts_io, ()), ('close_tasks', close_tasks_io, ()), ('close_proxy_access', close_proxy_access, ()),
        ('close_messages', close_message_io, ()), ('close_resources', close_resource_work, ()), ('close_open_sql', close_open_sql, ()), ('close_config_sql', close_config_sql, ()), ('close_shared', close_shared_work, ()),
        ('close_identity', close_identity_io, ()), ('close_auth', close_auth_io, ()),
        ('close_verification', close_verification_io, ()), ('close_profile', close_profile_io, ()),
        ('close_scripts', close_admin_scripts, ()), ('close_comments', close_comment_reads, ()),
        ('wait_periodic', wait_periodic_tasks, (workers,)),
    ):
        await failures.wait(stage, close, *args)
    app.state._ranking_shutdown_incomplete = failures.incomplete
    failures.raise_if_failed(cancel_requested=shutdown_cancel_requested(app.state))
    await run_owned_thread(db.flush_sync)
    logger.info("✅ HF 关闭前同步尝试已结束")
    logger.info("✅ 关闭完成")

# 🛡️ CORS 配置：默认通配符（插件前端运行于用户本地 ComfyUI 页面，origin 为
# 127.0.0.1/localhost/局域网 IP/自定义端口等无法穷举；本 API 通过 Authorization 头
# 传递 Token 且 allow_credentials=False 不依赖 Cookie，通配风险可控）。
# 如需收紧可在环境变量 ALLOWED_ORIGINS 中逗号分隔指定域名。
ALLOWED_ORIGINS = os.environ.get("ALLOWED_ORIGINS", "*").split(",")
logger.info(
    f"CORS 配置已加载: origin_count={len(ALLOWED_ORIGINS)} | "
    f"wildcard={len(ALLOWED_ORIGINS) == 1 and ALLOWED_ORIGINS[0] == '*'}"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=False, 
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["ETag", "X-Ranking-Error-Scope", "X-Ranking-Request-Id"],
)
app.add_middleware(RequestTraceMiddleware)

# ==========================================
# 路由挂载
# ==========================================
# 用户模块 (3个子模块)
app.include_router(users_auth_router)     # 🔐 登录/注册/密码重置
app.include_router(users_profile_router)  # 👤 用户资料
app.include_router(users_social_router)   # 🤝 关注/隐私

# 其他业务模块
app.include_router(items_router)          # 📦 内容管理
app.include_router(comments_router)       # 💬 评论系统
app.include_router(messages_router)       # ✉️ 私信系统
app.include_router(wallet_router)         # 💰 钱包/提现
app.include_router(proxy_router)          # 🔗 代理下载
app.include_router(posts_router)          # 💬 讨论区
app.include_router(prompts_router)        # ✨ 提示词市场
app.include_router(tasks_router)          # 📝 任务榜
app.include_router(open_api_router)       # 🌐 通用开放API（供第三方插件调用）
app.include_router(admin_plugins_router)  # 🔒 管理后台 - 第三方插件管理API

# ==========================================
# 🟢 私有图床代理中心 (Image Proxy)
# 解决 Private 仓库下，本地客户端报 401 Unauthorized 的终极方案
# ==========================================
from media_proxy_service import read_public_media
from media_file_response import OwnedMediaFileResponse

@app.get("/api/image_proxy")
@public_media_endpoint
async def proxy_hf_image(url: str = None, path: str = None):
    return await run_owned_thread(read_public_media, url, path, store=db, download=hf_hub_download,
                             file_response=OwnedMediaFileResponse, error_response=create_error_response)

# ==========================================
# 上传接口 (将返回的 URL 替换为 Proxy 代理链接)
# 🔒 P0安全优化：集成图片内容审核
# ==========================================
from image_moderation import moderate_image_sync, ModerationResult

# ==========================================
# 🎵 音频安全净化：重编码压缩，剥离元数据与隐藏附加信息
# ==========================================
from upload_service import process_upload

@app.post("/api/upload")
@shared_limit(30)
def upload_file(request: Request, file: UploadFile = File(...), file_type: str = Form(...), current_user: str = Depends(require_auth)):
    return process_upload(file, file_type, current_user, store=db, storage_factory=HfApi,
                          moderate=moderate_image_sync, limiter=_UPLOAD_SEMAPHORE,
                          error_response=create_error_response)



from github_source import select_github_token


class ValidateResourceRequest(BaseModel):
    url: str
    item_id: str
    account: str

from resource_probe_service import probe_resource
from github_probe_transport import open_github_probe
from skill_preflight import verify_skill_repository
from router_proxy import _send_with_validated_redirects
from safe_http_transport import ValidatedAsyncHTTPTransport
import httpx


def _probe_skill_repository(item, token):
    return verify_skill_repository(item, token, send=_send_with_validated_redirects,
        client_factory=httpx.AsyncClient, transport_factory=ValidatedAsyncHTTPTransport)

@app.post("/api/validate_resource")
@shared_limit(30)
def validate_resource(req_data: ValidateResourceRequest, request: Request, current_user: str = Depends(require_auth)):
    return probe_resource(req_data.item_id, store=db, download=hf_hub_download,
                          open_url=open_github_probe, token_selector=select_github_token,
                          error_response=create_error_response, skill_probe=_probe_skill_repository)

# ==========================================
# 🔒 管理员：系统配置 API
# ==========================================
app.include_router(system_config_router)

# ==========================================
# 🖥️ 管理界面静态文件挂载
# ⚠️ 必须放在所有 API 路由注册之后，避免拦截 API
# 访问 /admin/ 会自动返回 管理界面/index.html
# ==========================================
from fastapi.staticfiles import StaticFiles
app.mount("/admin", StaticFiles(directory="管理界面", html=True), name="admin")

# ==========================================
# 📱 移动端网页静态挂载（复用前端的 Web 版）
# 部署约定：同步 HF Space 时，除本目录内容外，
# 须将仓库根目录的 前端页面/ 整个目录一并复制到 Space 仓库根目录
# /m/assets 提供共享前端模块；/m 返回 移动端网页/index.html
# ⚠️ 必须放在所有 API 路由之后；先挂载更具体的 /m/assets，再挂载 /m
# 容错：目录缺失时仅跳过挂载并警告，不阻断 API 服务启动（移动端页面将 404）
# ==========================================
if os.path.isdir("前端页面"):
    app.mount("/m/assets", StaticFiles(directory="前端页面"), name="web_assets")
    app.mount("/m", StaticFiles(directory="移动端网页", html=True), name="mobile_web")
else:
    logger.warning("⚠️ 未找到 前端页面/ 目录，移动端网页 /m/ 未挂载（部署时须将仓库根目录的 前端页面/ 复制到 Space 仓库根目录），API 服务不受影响")