comfyui_o1key/__init__.py

"""
Comfyui_o1key - ComfyUI 自定义节点集合
通过 api.o1key.cn 调用 AI 模型进行图像生成和文本生成

项目结构:
├── nodes/          # 节点实现
├── utils/          # 工具模块
├── clients/        # API 客户端
└── __init__.py     # 节点注册入口
"""


import ssl
import logging
import asyncio

# 屏蔽 ComfyUI 资产扫描的终端日志输出
_seeder_filter = lambda record: not any(
    kw in record.getMessage()
    for kw in ("Seeder start", "Asset scan", "Scan(", "Fast scan")
)
logging.getLogger().addFilter(_seeder_filter)


def _is_ignored_asyncio_win10054(context):
    exc = context.get("exception")
    if not (
        isinstance(exc, ConnectionResetError)
        and getattr(exc, "winerror", None) == 10054
    ):
        return False

    handle = str(context.get("handle", ""))
    message = str(context.get("message", ""))
    marker = "_ProactorBasePipeTransport._call_connection_lost"
    return marker in handle or marker in message


def _install_asyncio_win10054_filter(loop):
    if getattr(loop, "_o1key_win10054_filter_installed", False):
        return loop

    previous_handler = loop.get_exception_handler()

    def _o1key_asyncio_exception_handler(loop, context):
        if _is_ignored_asyncio_win10054(context):
            return
        if previous_handler is not None:
            previous_handler(loop, context)
        else:
            loop.default_exception_handler(context)

    loop.set_exception_handler(_o1key_asyncio_exception_handler)
    setattr(loop, "_o1key_win10054_filter_installed", True)
    return loop


try:
    _install_asyncio_win10054_filter(asyncio.get_event_loop())
except RuntimeError:
    pass

if not getattr(asyncio, "_o1key_new_event_loop_patched", False):
    _o1key_original_new_event_loop = asyncio.new_event_loop

    def _o1key_new_event_loop(*args, **kwargs):
        return _install_asyncio_win10054_filter(
            _o1key_original_new_event_loop(*args, **kwargs)
        )

    asyncio.new_event_loop = _o1key_new_event_loop
    asyncio._o1key_new_event_loop_patched = True

from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, LoadImagesFromFolder, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, Google31Video, MiniMaxH3Video, FluxImageEdit, UniversalLLMChat, BatchImagesO1key, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, O1keyGrokVideo, O1keyGrokVideoEdit
from .nodes import K3Video, K3MotionControl, SaveImageFormat
from .nodes import O1keySavePSD
from .nodes import O1keyRemoveBackground
from .nodes import O1keyGridSplitter
from .nodes import O1keyPromptMultiFunction
from .nodes import O1keyVideoTrim
from .nodes import SeedanceElementCreate
from .nodes import SeedanceAutoPass
from .nodes import SeedanceAutoPassBatch
from .nodes import O1keyAutoRedCast
from .nodes import O1keyImageGenerator, O1keyImageSave
from .nodes import O1keyVideoGenerator, O1keyVideoResult
from .nodes import O1keyOmniFlashVideo

# 报错弹框友好文案（不修改原节点代码，仅在外层统一处理）
_MSG_TIMEOUT = "API 请求超时，请稍后重试或检查网络。"
_MSG_SSL_NETWORK = (
    "本地网络不太稳定!解决方案如下:\n"
    "1. 重启程序再试试看 (优先)\n"
    "2. 调整一下网络环境,如wifi或宽带等\n"
    "3. 切换VPN节点,或更换代理模式\n"
    "4. 关掉杀毒软件或防火墙\n"
    "5. 关掉浏览器VPN插件,避免冲突"
)

def _wrap_generate_for_error_display(cls, attr="generate"):
    original = getattr(cls, attr, None)
    if original is None:
        return
    def wrapped(self, *args, **kwargs):
        try:
            return original(self, *args, **kwargs)
        except TimeoutError as e:
            msg = (str(e) or "").strip()
            if not msg:
                msg = _MSG_TIMEOUT
            raise TimeoutError(msg) from None
        except (ssl.SSLError, OSError) as e:
            err_str = str(e)
            if "DECRYPTION_FAILED_OR_BAD_RECORD_MAC" in err_str or "decryption failed or bad record mac" in err_str.lower():
                raise RuntimeError(_MSG_SSL_NETWORK) from None
            raise
    setattr(cls, attr, wrapped)

_wrap_generate_for_error_display(NanoBananaPro)
_wrap_generate_for_error_display(BatchNanoBananaPro)

# ComfyUI 节点注册
NODE_CLASS_MAPPINGS = {
    "NanoBanana": NanoBananaPro,
    "BatchNanoBananaPro": BatchNanoBananaPro,
    "GoogleGemini": GoogleGemini,
    "LoadFile": LoadFile,
    "O1keyLoadImagesFromFolder": LoadImagesFromFolder,
    "ImageStitchPro": ImageStitchPro,

    "BatchCleanMetadata": BatchCleanMetadata,
    "VideoPreview": VideoPreview,
    "GoogleVeo": GoogleVeo,
    "Google31Video": Google31Video,
    "MiniMaxH3Video": MiniMaxH3Video,
    "FluxImageEdit": FluxImageEdit,
    "UniversalLLMChat": UniversalLLMChat,

    "BatchImagesO1key": BatchImagesO1key,
    "SeedanceMultiModal": SeedanceMultiModal,
    "StreamPreview": StreamPreview,
    "DoubaoImage": DoubaoImage,
    "O1keyGPTImage": O1keyGPTImage,
    "O1keyGPTImageBatch": O1keyGPTImageBatch,
    "O1keyGrokImage": O1keyGrokImage,
    "O1keyGrokVideo": O1keyGrokVideo,
    "O1keyGrokVideoEdit": O1keyGrokVideoEdit,
    "K3Video": K3Video,
    "K3MotionControl": K3MotionControl,
    "SaveImageFormat": SaveImageFormat,
    "O1keySavePSD": O1keySavePSD,
    "O1keyRemoveBackground": O1keyRemoveBackground,
    "O1keyGridSplitter": O1keyGridSplitter,
    "O1keyPromptMultiFunction": O1keyPromptMultiFunction,
    "O1keyVideoTrim": O1keyVideoTrim,
    "SeedanceElementCreate": SeedanceElementCreate,
    "SeedanceAutoPass": SeedanceAutoPass,
    "SeedanceAutoPassBatch": SeedanceAutoPassBatch,
    "O1keyAutoRedCast": O1keyAutoRedCast,
    "O1keyImageGenerator": O1keyImageGenerator,
    "O1keyImageSave": O1keyImageSave,
    "O1keyVideoGenerator": O1keyVideoGenerator,
    "O1keyVideoResult": O1keyVideoResult,
    "O1keyOmniFlashVideo": O1keyOmniFlashVideo,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "O1keyVideoGenerator": "o1key 视频生成",
    "O1keyVideoResult": "o1key 视频结果",
    "O1keyOmniFlashVideo": "Omni Flash 视频生成",
    "NanoBanana": "Nano Banana",
    "BatchNanoBananaPro": "Nano Banana 批量跑图",
    "GoogleGemini": "Google Gemini",
    "LoadFile": "加载文件",
    "O1keyLoadImagesFromFolder": "加载图像（文件夹）",
    "ImageStitchPro": "图像拼接 Pro",

    "BatchCleanMetadata": "批量任务（防AI识别）",
    "VideoPreview": "预览视频",
    "GoogleVeo": "Google Veo - ab",
    "Google31Video": "Google 3.1 Video",
    "MiniMaxH3Video": "MiniMax H3 / H3 Max 视频生成",
    "FluxImageEdit": "Flux2 图像编辑",
    "UniversalLLMChat": "提示词专家",

    "BatchImagesO1key": "加载图像（批量）",
    "SeedanceMultiModal": "Seedance 多模态参考生视频",
    "StreamPreview": "流式文本预览",
    "DoubaoImage": "豆包生图",
    "O1keyGPTImage": "gpt image",
    "O1keyGPTImageBatch": "o1key GPT Image（批量）",
    "O1keyGrokImage": "Grok Image",
    "O1keyGrokVideo": "Grok Video",
    "O1keyGrokVideoEdit": "Grok Video Edit",
    "K3Video": "K 视频生成",
    "K3MotionControl": "K 动作模仿",
    "SaveImageFormat": "保存图像（格式转换）",
    "O1keySavePSD": "保存 PSD（分层）",
    "O1keyRemoveBackground": "去背景（rembg）",
    "O1keyGridSplitter": "合并图智能切割",
    "O1keyPromptMultiFunction": "提示词（多功能）",
    "O1keyVideoTrim": "视频裁剪",
    "SeedanceElementCreate": "Seedance 创建素材",
    "SeedanceAutoPass": "Seedance 全能生成视频",
    "SeedanceAutoPassBatch": "Seedance 全能生成视频（批量）",
    "O1keyAutoRedCast": "自动红偏校正",
    "O1keyImageGenerator": "o1key 图片生成",
    "O1keyImageSave": "o1key 保存图像",
}

WEB_DIRECTORY = "./web"

__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']

# 注册 /o1key/input_dir 接口，供前端文件上传按钮获取 input 目录绝对路径
try:
    from aiohttp import web
    from server import PromptServer
    import folder_paths
    from .utils.config import (
        DEFAULT_NETWORK_ROUTE,
        NETWORK_ROUTE_CONFIG_KEY,
        NETWORK_ROUTE_OPTIONS,
        NETWORK_ROUTES,
        get_api_key,
        get_network_route,
        load_config,
        update_config,
    )
    from .utils.chat_support import (
        PROMPT_OPTIMIZER_TIMEOUT_SECONDS,
        build_search_context,
        expand_xlsx_attachments,
        extract_search_query,
        optimize_image_prompt,
        rewrite_search_query,
        web_search,
        write_video_prompt,
    )
    from .utils.o1key_image_jobs import register_o1key_image_job_routes
    from .utils.o1key_image_thumbnail import register_o1key_image_thumbnail_route
    from .utils.o1key_video_jobs import register_o1key_video_job_routes
    from .utils.updater import UpdateError, update_package

    _O1KEY_IMAGE_JOB_MANAGER = register_o1key_image_job_routes(
        PromptServer,
        web,
        folder_paths,
    )
    register_o1key_image_thumbnail_route(PromptServer, web, folder_paths)
    _O1KEY_VIDEO_JOB_MANAGER = register_o1key_video_job_routes(
        PromptServer,
        web,
        folder_paths,
    )

    # 每次 ComfyUI 进程启动都会生成新的标识。前端据此确认后端确实完成了
    # 重启，而不是仅仅重新加载了浏览器页面。
    import os as _restart_os
    import sys as _restart_sys
    import threading as _restart_threading
    import time as _restart_time
    import uuid as _restart_uuid

    _O1KEY_BOOT_ID = _restart_uuid.uuid4().hex
    _o1key_restart_pending = _restart_threading.Event()
    _o1key_update_lock = _restart_threading.Lock()

    @PromptServer.instance.routes.post("/o1key/update")
    async def update_o1key_package(request):
        if request.headers.get("X-O1Key-Update") != "1":
            return web.json_response(
                {"code": "invalid_request", "error": "无效的更新请求。", "suggestion": "请从 O1Key 更新面板重新操作。"},
                status=403,
            )
        if not _o1key_update_lock.acquire(blocking=False):
            return web.json_response(
                {"code": "update_in_progress", "error": "更新正在进行。", "suggestion": "请等待当前操作完成，不要重复点击。"},
                status=409,
            )
        try:
            result = await asyncio.to_thread(update_package)
            return web.json_response(result)
        except UpdateError as exc:
            return web.json_response(exc.as_dict(), status=exc.status)
        except Exception:
            logging.exception("O1Key 更新失败")
            return web.json_response(
                {"code": "internal_error", "error": "更新失败。", "suggestion": "请查看 ComfyUI 日志，并在确认本地文件安全后重试。"},
                status=500,
            )
        finally:
            _o1key_update_lock.release()

    def _o1key_restart_command():
        """复用当前解释器和启动参数，并禁止重启时额外打开浏览器。"""
        auto_launch_flags = {"--auto-launch", "--auto_launch", "--launch"}
        # orig_argv 包含嵌入式 Python 的 -s 等解释器参数；普通 sys.argv 不包含。
        # 保留这些参数可确保便携版重启前后的运行环境完全一致。
        original = getattr(_restart_sys, "orig_argv", None)
        source_arguments = original[1:] if original else _restart_sys.argv
        arguments = [arg for arg in source_arguments if arg not in auto_launch_flags]
        if "--disable-auto-launch" not in arguments:
            arguments.append("--disable-auto-launch")
        return [_restart_sys.executable, *arguments]

    def _restart_o1key_comfyui_process(delay=1.25):
        """在响应发送完成后，用相同终端进程重新启动 ComfyUI。"""
        try:
            _restart_time.sleep(delay)
            try:
                _restart_sys.stdout.flush()
                _restart_sys.stderr.flush()
            except Exception:
                pass

            command = _o1key_restart_command()
            print("[O1Key] 正在重启 ComfyUI...", flush=True)
            _restart_os.execv(command[0], command)
        except Exception:
            _o1key_restart_pending.clear()
            logging.exception("O1Key 无法重启 ComfyUI 进程")

    @PromptServer.instance.routes.get("/o1key/restart/status")
    async def get_o1key_restart_status(request):
        return web.json_response(
            {
                "ready": True,
                "boot_id": _O1KEY_BOOT_ID,
                "pid": _restart_os.getpid(),
            },
            headers={"Cache-Control": "no-store"},
        )

    @PromptServer.instance.routes.post("/o1key/restart")
    async def restart_o1key_comfyui(request):
        if _o1key_restart_pending.is_set():
            return web.json_response(
                {
                    "success": False,
                    "error": "ComfyUI 正在重启，请稍候。",
                    "boot_id": _O1KEY_BOOT_ID,
                },
                status=409,
            )

        _o1key_restart_pending.set()
        worker = _restart_threading.Thread(
            target=_restart_o1key_comfyui_process,
            name="o1key-comfyui-restart",
            daemon=True,
        )
        worker.start()
        return web.json_response(
            {
                "success": True,
                "message": "ComfyUI 正在重启。",
                "boot_id": _O1KEY_BOOT_ID,
                "pid": _restart_os.getpid(),
            },
            headers={"Cache-Control": "no-store"},
        )

    def _get_o1key_server_port():
        try:
            import comfy.cli_args as _cli_args
            args = getattr(_cli_args, "args", None)
            port = getattr(args, "port", None) if args else None
            port = port or getattr(_cli_args, "server_port", None) or getattr(_cli_args, "port", None)
            if port is not None:
                return str(int(port))
        except Exception:
            pass
        try:
            import sys as _sys
            for idx, arg in enumerate(_sys.argv):
                if arg in ("--port", "--listen-port") and idx + 1 < len(_sys.argv):
                    return str(int(_sys.argv[idx + 1]))
                for prefix in ("--port=", "--listen-port="):
                    if arg.startswith(prefix):
                        return str(int(arg.split("=", 1)[1]))
        except Exception:
            pass
        return "8188"

    def _get_o1key_history_meta_file(output_dir):
        import os as _os_history
        return _os_history.path.join(
            output_dir,
            f".o1key_history_{_get_o1key_server_port()}.json",
        )

    def _get_o1key_notes_file():
        # 笔记固定存 ComfyUI input 目录，插件更新/替换不会清空笔记
        import os as _os_notes
        input_dir = _os_notes.path.abspath(folder_paths.get_input_directory())
        _os_notes.makedirs(input_dir, exist_ok=True)
        notes_file = _os_notes.path.join(input_dir, "o1key-notes.json")

        return notes_file

    def _get_o1key_cases_dir():
        import os as _os_cases
        cases_dir = _os_cases.path.join(_os_cases.path.dirname(__file__), "cases")
        _os_cases.makedirs(cases_dir, exist_ok=True)
        return cases_dir

    def _get_o1key_case_file(filename):
        import os as _os_cases
        safe_name = _os_cases.path.basename(filename or "")
        if not safe_name.lower().endswith(".json"):
            return None
        cases_dir = _get_o1key_cases_dir()
        path = _os_cases.path.abspath(_os_cases.path.join(cases_dir, safe_name))
        if not path.startswith(_os_cases.path.abspath(cases_dir) + _os_cases.sep):
            return None
        return path

    def _extract_o1key_notes(payload):
        if isinstance(payload, list):
            return payload
        if isinstance(payload, dict) and isinstance(payload.get("notes"), list):
            return payload["notes"]
        return None

    @PromptServer.instance.routes.get("/o1key/cases")
    async def get_o1key_cases(request):
        import os as _os_cases
        import json as _json_cases

        cases_dir = _get_o1key_cases_dir()
        cases = []
        for filename in sorted(_os_cases.listdir(cases_dir), key=str.lower):
            if not filename.lower().endswith(".json"):
                continue
            path = _get_o1key_case_file(filename)
            if not path or not _os_cases.path.isfile(path):
                continue
            title = _os_cases.path.splitext(filename)[0]
            try:
                with open(path, "r", encoding="utf-8") as cf:
                    data = _json_cases.load(cf)
                if isinstance(data, dict):
                    title = str(data.get("title") or data.get("name") or title)
            except Exception:
                pass
            cases.append({"id": filename, "filename": filename, "title": title})
        return web.json_response({"cases": cases, "path": str(cases_dir)})

    @PromptServer.instance.routes.get("/o1key/case")
    async def get_o1key_case(request):
        import os as _os_cases
        import json as _json_cases

        filename = request.query.get("file", "")
        path = _get_o1key_case_file(filename)
        if not path or not _os_cases.path.isfile(path):
            return web.json_response({"error": "case not found"}, status=404)
        try:
            with open(path, "r", encoding="utf-8") as cf:
                data = _json_cases.load(cf)
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)
        return web.json_response({"filename": _os_cases.path.basename(path), "case": data})

    @PromptServer.instance.routes.get("/o1key/notes")
    async def get_o1key_notes(request):
        import os as _os_notes
        import json as _json_notes

        notes_file = _get_o1key_notes_file()
        exists = _os_notes.path.isfile(notes_file)
        notes = []

        if exists:
            try:
                with open(notes_file, "r", encoding="utf-8") as nf:
                    loaded = _json_notes.load(nf)
                notes = _extract_o1key_notes(loaded)
                if notes is None:
                    return web.json_response(
                        {"error": "invalid notes file", "path": notes_file},
                        status=500,
                    )
            except Exception as e:
                return web.json_response(
                    {"error": str(e), "path": notes_file},
                    status=500,
                )

        return web.json_response({"notes": notes, "path": notes_file, "exists": exists})

    @PromptServer.instance.routes.post("/o1key/notes")
    async def save_o1key_notes(request):
        import os as _os_notes
        import json as _json_notes

        try:
            payload = await request.json()
            notes = _extract_o1key_notes(payload)
            if notes is None:
                return web.json_response({"error": "notes must be a list"}, status=400)
        except Exception as e:
            return web.json_response({"error": f"invalid notes payload: {str(e)}"}, status=400)

        notes_file = _get_o1key_notes_file()
        temp_file = notes_file + ".tmp"
        try:
            with open(temp_file, "w", encoding="utf-8") as nf:
                _json_notes.dump(notes, nf, ensure_ascii=False, indent=2)
                nf.write("\n")
            _os_notes.replace(temp_file, notes_file)
        except Exception as e:
            return web.json_response({"error": f"save notes failed: {str(e)}"}, status=500)

        return web.json_response({
            "success": True,
            "path": notes_file,
            "count": len(notes),
        })

    @PromptServer.instance.routes.get("/o1key/input_dir")
    async def get_input_dir(request):
        import os
        path = os.path.abspath(folder_paths.get_input_directory())
        return web.json_response({"path": path})

    @PromptServer.instance.routes.get("/o1key/api_key")
    async def get_api_key_route(request):
        config = load_config()
        key = config.get("O1KEY_API_KEY", "")
        masked = ""
        if key:
            if len(key) > 8:
                masked = key[:3] + "****" + key[-4:]
            else:
                masked = "****"
        return web.json_response({
            "has_key": bool(key),
            "masked": masked,
            "network_route": get_network_route(),
            "network_route_options": NETWORK_ROUTE_OPTIONS,
        })

    @PromptServer.instance.routes.post("/o1key/config")
    async def set_o1key_config(request):
        data = await request.json()
        route = str(data.get("network_route", "")).strip()
        if route not in NETWORK_ROUTES:
            return web.json_response({"error": "网络线路无效"}, status=400)

        updates = {NETWORK_ROUTE_CONFIG_KEY: route}
        if data.get("api_key") is not None:
            new_key = str(data.get("api_key", "")).strip()
            if not new_key:
                return web.json_response({"error": "API Key 不能为空"}, status=400)
            if "\n" in new_key or "\r" in new_key:
                return web.json_response({"error": "API Key 格式无效"}, status=400)
            updates["O1KEY_API_KEY"] = new_key

        config = update_config(updates=updates)
        key = config.get("O1KEY_API_KEY", "")
        masked = key[:3] + "****" + key[-4:] if len(key) > 8 else ("****" if key else "")
        return web.json_response({
            "success": True,
            "has_key": bool(key),
            "masked": masked,
            "network_route": route,
        })

    @PromptServer.instance.routes.post("/o1key/network_route")
    async def set_network_route(request):
        data = await request.json()
        route = str(data.get("network_route", "")).strip()
        if route not in NETWORK_ROUTES:
            return web.json_response({"error": "网络线路无效"}, status=400)
        update_config(updates={NETWORK_ROUTE_CONFIG_KEY: route})
        return web.json_response({"success": True, "network_route": route})

    @PromptServer.instance.routes.post("/o1key/api_key")
    async def set_api_key_route(request):
        data = await request.json()
        new_key = data.get("api_key", "").strip()
        if not new_key:
            return web.json_response({"error": "API Key 不能为空"}, status=400)
        if "\n" in new_key or "\r" in new_key:
            return web.json_response({"error": "API Key 格式无效"}, status=400)
        update_config(updates={"O1KEY_API_KEY": new_key})
        return web.json_response({"success": True})

    @PromptServer.instance.routes.post("/o1key/test_key")
    async def test_api_key_route(request):
        import aiohttp as _aiohttp
        data = await request.json()
        test_key = str(data.get("api_key") or "").strip() or get_api_key()
        if not test_key:
            return web.json_response({"valid": False, "error": "请先输入或保存 API Key"})
        requested_route = str(data.get("network_route", "")).strip()
        route = requested_route if requested_route in NETWORK_ROUTES else get_network_route()
        base_url = NETWORK_ROUTES.get(route, NETWORK_ROUTES[DEFAULT_NETWORK_ROUTE])
        url = f"{base_url}/v1/models"
        headers = {"Authorization": f"Bearer {test_key}"}
        try:
            async with _aiohttp.ClientSession() as session:
                async with session.get(url, headers=headers, timeout=_aiohttp.ClientTimeout(total=10)) as resp:
                    if resp.status == 200:
                        return web.json_response({"valid": True})
                    elif resp.status == 401:
                        return web.json_response({"valid": False, "error": "密钥无效或已过期"})
                    else:
                        text = await resp.text()
                        return web.json_response({"valid": False, "error": f"验证失败 ({resp.status})"})
        except Exception as e:
            return web.json_response({"valid": False, "error": f"网络错误: {str(e)}"})

    @PromptServer.instance.routes.delete("/o1key/api_key")
    async def delete_api_key_route(request):
        update_config(remove=["O1KEY_API_KEY"])
        return web.json_response({"success": True})

    # === 主体（Element）代理：转发到 {base}/kling/v1/general/*，后端注入令牌 ===
    _ELEMENT_PREFIX = "/kling/v1/general"

    def _element_base_url(_route=None):
        return NETWORK_ROUTES[get_network_route()].rstrip("/")

    def _element_headers():
        config = load_config()
        key = config.get("O1KEY_API_KEY", "")
        if not key:
            return None
        return {"Authorization": f"Bearer {key}"}

    @PromptServer.instance.routes.get("/o1key/element/mine")
    async def o1key_element_mine(request):
        """列表接口：GET /kling/v1/general/advanced-custom-elements"""
        import aiohttp as _aiohttp
        headers = _element_headers()
        if not headers:
            return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
        base = _element_base_url(request.query.get("route"))
        url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements"
        params = {}
        # 支持分页参数
        page_num = request.query.get("pageNum", "1")
        page_size = request.query.get("pageSize", "100")
        params["pageNum"] = page_num
        params["pageSize"] = page_size
        try:
            async with _aiohttp.ClientSession() as session:
                async with session.get(url, headers=headers, params=params,
                                       timeout=_aiohttp.ClientTimeout(total=30)) as up:
                    result = await up.json()
                    # 新API返回: {"success": true, "data": {"code": 0, "data": [...], "total": N}, "message": ""}
                    # 转换为前端期望的格式: {"success": true, "data": [...]}
                    if result.get("success") and isinstance(result.get("data"), dict):
                        elements = result["data"].get("data", [])
                        return web.json_response({"success": True, "data": elements, "message": ""})
                    return web.json_response(result, status=up.status)
        except Exception as e:
            return web.json_response({"success": False, "message": str(e)}, status=502)

    @PromptServer.instance.routes.post("/o1key/element/upload")
    async def o1key_element_upload(request):
        """转发 multipart 文件上传：POST /kling/v1/general/upload
        视频可达 200MB，固定总超时会截断大文件上传，改用：不限总时长 +
        读空闲 120s 超时（连接卡死才超时，慢速大文件不会被一刀切断）。"""
        import aiohttp as _aiohttp
        headers = _element_headers()
        if not headers:
            return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
        base = _element_base_url(request.query.get("route"))
        url = f"{base}{_ELEMENT_PREFIX}/upload"
        try:
            reader = await request.multipart()
            field = await reader.next()
            if field is None or field.name != "file":
                return web.json_response({"success": False, "message": "缺少 file 字段"}, status=400)
            file_bytes = await field.read(decode=False)
            filename = field.filename or "image.png"
            form = _aiohttp.FormData()
            form.add_field("file", file_bytes, filename=filename,
                           content_type=field.headers.get("Content-Type", "application/octet-stream"))
            timeout = _aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=120)
            async with _aiohttp.ClientSession() as session:
                async with session.post(url, headers=headers, data=form,
                                        timeout=timeout) as up:
                    data = await up.json()
                    return web.json_response(data, status=up.status)
        except Exception as e:
            return web.json_response({"success": False, "message": str(e)}, status=502)

    @PromptServer.instance.routes.post("/o1key/element/create")
    async def o1key_element_create(request):
        """创建主体：POST /kling/v1/general/advanced-custom-elements

        新API字段映射：
        - name -> element_name
        - description -> element_description
        - reference_type -> reference_type (image_refer / video_refer)
        - frontal_image -> frontal_image
        - refer_images -> refer_images
        - video_list -> video_list
        - element_voice_id, tag_ids, channel_id 保持不变
        """
        import aiohttp as _aiohttp
        import json as _json_element
        headers = _element_headers()
        if not headers:
            return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
        try:
            payload = await request.json()
        except Exception as e:
            return web.json_response({"success": False, "message": f"请求体无效: {e}"}, status=400)
        route = payload.pop("route", None)
        base = _element_base_url(route)
        url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements"

        # 字段映射：前端使用旧字段名，转换为新API字段名
        api_payload = {}
        if "name" in payload:
            api_payload["element_name"] = payload["name"]
        if "description" in payload:
            api_payload["element_description"] = payload["description"]
        # 其他字段直接透传
        for key in ["reference_type", "frontal_image", "refer_images", "video_list",
                    "element_voice_id", "tag_ids", "channel_id"]:
            if key in payload:
                api_payload[key] = payload[key]

        send_headers = {**headers, "Content-Type": "application/json"}
        # 打印创建主体的请求信息
        try:
            print(f"[主体创建] 请求 URL: {url}")
            print("[主体创建] 请求体: " + _json_element.dumps(api_payload, ensure_ascii=False, indent=2))
        except Exception:
            pass
        try:
            async with _aiohttp.ClientSession() as session:
                async with session.post(url, headers=send_headers, json=api_payload,
                                        timeout=_aiohttp.ClientTimeout(total=60)) as up:
                    data = await up.json()
                    # 打印创建主体的响应信息
                    try:
                        print("[主体创建] 响应体: " + _json_element.dumps(data, ensure_ascii=False, indent=2))
                    except Exception:
                        pass
                    return web.json_response(data, status=up.status)
        except Exception as e:
            return web.json_response({"success": False, "message": str(e)}, status=502)

    @PromptServer.instance.routes.post("/o1key/element/refresh")
    async def o1key_element_refresh(request):
        """查询主体：GET /kling/v1/general/advanced-custom-elements/{task_id}

        前端传 id（数据库主键），需要先查本地库拿到 job_id（即 task_id），再查询上游。
        为了简化，这里改为前端直接传 task_id（即创建时返回的 job_id）。
        """
        import aiohttp as _aiohttp
        import json as _json_refresh
        headers = _element_headers()
        if not headers:
            return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
        try:
            payload = await request.json()
        except Exception:
            payload = {}
        task_id = payload.get("task_id") or payload.get("id")
        if not task_id:
            return web.json_response({"success": False, "message": "缺少 task_id"}, status=400)
        base = _element_base_url(payload.get("route"))
        url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements/{task_id}"
        try:
            async with _aiohttp.ClientSession() as session:
                async with session.get(url, headers=headers,
                                       timeout=_aiohttp.ClientTimeout(total=60)) as up:
                    data = await up.json()
                    # 打印查询响应
                    try:
                        print(f"[主体查询] task_id={task_id}")
                        print("[主体查询] 响应体: " + _json_refresh.dumps(data, ensure_ascii=False, indent=2))
                    except Exception:
                        pass
                    # 新API返回嵌套结构，需要提取 task_status 和 element_id
                    # 响应: {"success": true, "data": {"code": 0, "data": {"task_status": "succeed", "task_result": {"elements": [...]}}}}
                    if data.get("success") and isinstance(data.get("data"), dict):
                        inner = data["data"].get("data", {})
                        task_status = inner.get("task_status", "")
                        # 转换为前端期望的格式
                        element = {
                            "id": task_id,
                            "job_id": task_id,
                            "status": task_status,
                            "task_status": task_status,
                        }
                        if task_status == "succeed":
                            elements = inner.get("task_result", {}).get("elements", [])
                            if elements:
                                first = elements[0]
                                element["element_id"] = str(first.get("element_id", ""))
                                element["name"] = first.get("element_name", "")
                                element["frontal_image"] = first.get("element_image_list", {}).get("frontal_image", "")
                        elif task_status == "failed":
                            element["fail_reason"] = inner.get("task_status_msg", "")
                        return web.json_response({"success": True, "data": {"element": element, "detail": data["data"]}})
                    return web.json_response(data, status=up.status)
        except Exception as e:
            return web.json_response({"success": False, "message": str(e)}, status=502)

    @PromptServer.instance.routes.post("/o1key/element/delete")
    async def o1key_element_delete(request):
        """删除主体：POST /kling/v1/general/delete-advanced-elements

        请求体: {"element_id": "315320838184520"}
        """
        import aiohttp as _aiohttp
        headers = _element_headers()
        if not headers:
            return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
        try:
            payload = await request.json()
        except Exception:
            payload = {}
        element_id = payload.get("element_id") or payload.get("id")
        if not element_id:
            return web.json_response({"success": False, "message": "缺少 element_id"}, status=400)
        base = _element_base_url(payload.get("route"))
        url = f"{base}{_ELEMENT_PREFIX}/delete-advanced-elements"
        delete_payload = {"element_id": str(element_id)}
        send_headers = {**headers, "Content-Type": "application/json"}
        try:
            async with _aiohttp.ClientSession() as session:
                async with session.post(url, headers=send_headers, json=delete_payload,
                                       timeout=_aiohttp.ClientTimeout(total=60)) as up:
                    data = await up.json()
                    return web.json_response(data, status=up.status)
        except Exception as e:
            return web.json_response({"success": False, "message": str(e)}, status=502)

    @PromptServer.instance.routes.get("/o1key/element/image")
    async def o1key_element_image(request):
        """图片同源代理：ComfyUI 的 CSP 限制 img-src 'self'，外链缩略图无法直接显示。
        前端把缩略图 src 指向本路由，后端取回字节再吐给浏览器，对浏览器即同源。
        仅允许 o1key 资源域，避免被当成任意 URL 抓取的 SSRF 跳板。"""
        import aiohttp as _aiohttp
        from urllib.parse import urlparse, unquote
        raw = request.query.get("url", "")
        if not raw:
            return web.json_response({"success": False, "message": "缺少 url"}, status=400)
        target = unquote(raw)
        try:
            parsed = urlparse(target)
        except Exception:
            parsed = None
        if not parsed or parsed.scheme not in ("http", "https"):
            return web.json_response({"success": False, "message": "非法 url"}, status=400)
        host = (parsed.hostname or "").lower()
        if not (host.endswith(".o1key.com") or host.endswith(".o1key.cn")
                or host in ("o1key.com", "o1key.cn")):
            return web.json_response({"success": False, "message": "不允许的图片来源"}, status=403)
        try:
            async with _aiohttp.ClientSession() as session:
                async with session.get(target, timeout=_aiohttp.ClientTimeout(total=30)) as up:
                    if up.status != 200:
                        return web.Response(status=up.status)
                    body = await up.read()
                    ctype = up.headers.get("Content-Type", "image/jpeg").split(";")[0].strip()
            return web.Response(body=body, content_type=ctype or "image/jpeg",
                                headers={"Cache-Control": "max-age=3600"})
        except Exception as e:
            return web.json_response({"success": False, "message": str(e)}, status=502)

    @PromptServer.instance.routes.get("/o1key/output_history")
    async def get_output_history(request):
        """读取 output 目录文件，按执行分组返回 /api/jobs 兼容格式"""
        import os, json as _json
        limit = int(request.query.get("limit", "200"))
        offset = int(request.query.get("offset", "0"))
        output_dir = os.path.abspath(folder_paths.get_output_directory())
        meta_file = _get_o1key_history_meta_file(output_dir)
        meta = {}
        if os.path.isfile(meta_file):
            try:
                with open(meta_file, "r", encoding="utf-8") as mf:
                    meta = _json.load(mf)
            except Exception:
                pass
        supported_ext = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.webm'}
        # 收集所有文件并按 workflow_id 分组
        all_files = []
        for fname in meta.keys():
            ext = os.path.splitext(fname)[1].lower()
            if ext not in supported_ext:
                continue
            fpath = os.path.join(output_dir, fname)
            if not os.path.isfile(fpath):
                continue
            mtime = os.path.getmtime(fpath)
            media = "images" if ext in {'.png','.jpg','.jpeg','.webp','.gif'} else "video"
            all_files.append({"name": fname, "mtime": mtime, "media": media})
        # 按 workflow_id 分组（同一次执行合并为一个 job）
        groups = {}
        for f in all_files:
            m = meta.get(f["name"], {})
            wid = m.get("workflow_id")
            if wid:
                groups.setdefault(wid, []).append((f, m))
        # 构建 job 列表
        jobs = []
        for wid, items in groups.items():
            items.sort(key=lambda x: x[0]["mtime"], reverse=True)
            latest = items[0]
            f, m = latest
            start_ms = int(m.get("start_time", f["mtime"]) * 1000)
            end_ms = int(m.get("end_time", f["mtime"]) * 1000)
            jobs.append({
                "id": wid,
                "status": "completed",
                "create_time": start_ms,
                "execution_start_time": start_ms,
                "execution_end_time": end_ms,
                "preview_output": {
                    "filename": f["name"],
                    "subfolder": "",
                    "type": "output",
                    "nodeId": "0",
                    "mediaType": f["media"],
                },
                "outputs_count": len(items),
                "execution_error": None,
                "workflow_id": wid,
            })
        # 按时间倒序排列，分页
        jobs.sort(key=lambda x: x["create_time"], reverse=True)
        total = len(jobs)
        page = jobs[offset:offset+limit]
        return web.json_response({
            "jobs": page,
            "pagination": {"offset": offset, "limit": limit, "total": total, "has_more": offset + limit < total}
        })

    @PromptServer.instance.routes.get("/o1key/output_workflow")
    async def get_output_workflow(request):
        """从 PNG 元数据中读取工作流，供前端恢复使用"""
        import os, struct, json as _json
        filename = request.query.get("filename", "")
        if not filename:
            return web.json_response({"error": "missing filename"}, status=400)
        output_dir = os.path.abspath(folder_paths.get_output_directory())
        fpath = os.path.join(output_dir, filename)
        if not os.path.isfile(fpath) or not fpath.lower().endswith(".png"):
            return web.json_response({"error": "file not found"}, status=404)
        workflow = None
        prompt_data = None
        try:
            with open(fpath, "rb") as pf:
                pf.read(8)  # PNG signature
                while True:
                    raw = pf.read(8)
                    if len(raw) < 8:
                        break
                    length = struct.unpack(">I", raw[:4])[0]
                    chunk_type = raw[4:8]
                    data = pf.read(length)
                    pf.read(4)  # CRC
                    if chunk_type == b"tEXt":
                        key, val = data.split(b"\x00", 1)
                        k = key.decode("ascii", errors="replace")
                        if k == "workflow":
                            workflow = _json.loads(val)
                        elif k == "prompt":
                            prompt_data = _json.loads(val)
                    elif chunk_type == b"IEND":
                        break
        except Exception:
            pass
        return web.json_response({"workflow": workflow, "prompt": prompt_data})

    @PromptServer.instance.routes.get("/o1key/job_detail/{job_id}")
    async def get_job_detail(request):
        """根据 job_id 返回当前端口持久化历史中的 job 详情"""
        import os, struct, json as _json
        job_id = request.match_info["job_id"]
        output_dir = os.path.abspath(folder_paths.get_output_directory())
        meta_file = _get_o1key_history_meta_file(output_dir)
        meta = {}
        if os.path.isfile(meta_file):
            try:
                with open(meta_file, "r", encoding="utf-8") as mf:
                    meta = _json.load(mf)
            except Exception:
                pass
        supported_ext = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.webm'}
        # 只在当前端口的持久化记录中查找该 job 的文件
        matched_files = []
        for fname in meta.keys():
            ext = os.path.splitext(fname)[1].lower()
            if ext not in supported_ext:
                continue
            m = meta.get(fname, {})
            if m.get("workflow_id") == job_id:
                matched_files.append(fname)
        if not matched_files:
            return web.json_response({"error": "not found"}, status=404)
        # 用最新文件作为代表
        matched_files.sort(key=lambda f: os.path.getmtime(os.path.join(output_dir, f)), reverse=True)
        target_file = matched_files[0]
        fpath = os.path.join(output_dir, target_file)
        m = meta.get(target_file, {})
        mtime = os.path.getmtime(fpath)
        start_ms = int(m.get("start_time", mtime) * 1000)
        end_ms = int(m.get("end_time", mtime) * 1000)
        ext = os.path.splitext(target_file)[1].lower()
        media = "images" if ext in {'.png','.jpg','.jpeg','.webp','.gif'} else "video"
        # 读取 PNG 工作流元数据
        workflow = None
        if ext == ".png":
            try:
                with open(fpath, "rb") as pf:
                    pf.read(8)
                    while True:
                        raw = pf.read(8)
                        if len(raw) < 8:
                            break
                        length = struct.unpack(">I", raw[:4])[0]
                        chunk_type = raw[4:8]
                        data = pf.read(length)
                        pf.read(4)
                        if chunk_type == b"tEXt":
                            key, val = data.split(b"\x00", 1)
                            k = key.decode("ascii", errors="replace")
                            if k == "workflow":
                                workflow = _json.loads(val)
                        elif chunk_type == b"IEND":
                            break
            except Exception:
                pass
        # 构建 outputs：包含该执行的所有文件
        outputs = {}
        for i, fname in enumerate(matched_files):
            e = os.path.splitext(fname)[1].lower()
            mt = "images" if e in {'.png','.jpg','.jpeg','.webp','.gif'} else "gifs"
            outputs.setdefault(str(i), {}).setdefault(mt, []).append(
                {"filename": fname, "subfolder": "", "type": "output"}
            )
        job_detail = {
            "id": job_id,
            "status": "completed",
            "create_time": start_ms,
            "execution_start_time": start_ms,
            "execution_end_time": end_ms,
            "preview_output": {
                "filename": target_file,
                "subfolder": "",
                "type": "output",
                "nodeId": "0",
                "mediaType": media,
            },
            "outputs_count": len(matched_files),
            "execution_error": None,
            "workflow_id": job_id,
            "workflow": {
                "extra_data": {
                    "extra_pnginfo": {"workflow": workflow}
                }
            } if workflow else None,
            "outputs": outputs,
        }
        return web.json_response(job_detail)

    @PromptServer.instance.routes.post("/o1key/delete_history")
    async def delete_history_item(request):
        """删除持久化历史记录及对应的输出文件"""
        import os, json as _json
        body = await request.json()
        job_ids = body.get("delete", [])
        if not job_ids:
            return web.json_response({"success": False, "error": "missing ids"}, status=400)
        output_dir = os.path.abspath(folder_paths.get_output_directory())
        meta_file = _get_o1key_history_meta_file(output_dir)
        meta = {}
        if os.path.isfile(meta_file):
            try:
                with open(meta_file, "r", encoding="utf-8") as mf:
                    meta = _json.load(mf)
            except Exception:
                pass
        deleted_files = []
        for job_id in job_ids:
            files_to_remove = []
            for fname, m in list(meta.items()):
                if m.get("workflow_id") == job_id:
                    files_to_remove.append(fname)
            for fname in files_to_remove:
                meta.pop(fname, None)
                fpath = os.path.join(output_dir, fname)
                if os.path.isfile(fpath):
                    try:
                        os.remove(fpath)
                        deleted_files.append(fname)
                    except Exception:
                        pass
        try:
            with open(meta_file, "w", encoding="utf-8") as mf:
                _json.dump(meta, mf, ensure_ascii=False)
        except Exception:
            pass
        return web.json_response({"success": True, "deleted": deleted_files})

    # === 图片生成提示词优化（服务端读取参考图，避免前端接触 API Key） ===
    @PromptServer.instance.routes.post("/o1key/image/prompt-optimize")
    async def optimize_o1key_image_prompt(request):
        import aiohttp as _aiohttp

        try:
            data = await request.json()
            if not isinstance(data, dict):
                raise ValueError("请求体必须是对象")
            prompt = data.get("prompt", "")
            references = data.get("references", [])
            api_key = get_api_key() or ""
            if not api_key:
                return web.json_response({"error": "未配置 API Key"}, status=401)

            timeout = _aiohttp.ClientTimeout(total=PROMPT_OPTIMIZER_TIMEOUT_SECONDS + 15)
            async with _aiohttp.ClientSession(timeout=timeout) as session:
                optimized = await optimize_image_prompt(
                    session,
                    NETWORK_ROUTES[get_network_route()],
                    api_key,
                    prompt,
                    references,
                    folder_paths.get_input_directory(),
                )
            return web.json_response(
                {
                    "prompt": optimized,
                    "model": "gpt-5.6-sol",
                    "reasoning_effort": "high",
                },
                headers={"Cache-Control": "no-store"},
            )
        except ValueError as exc:
            return web.json_response({"error": str(exc)}, status=400)
        except RuntimeError as exc:
            return web.json_response({"error": str(exc)}, status=502)
        except Exception:
            return web.json_response({"error": "提示词优化失败，请稍后重试"}, status=500)

    # === 视频生成 AI帮写（独立视频预设，仅分析安全的 input 图片描述） ===
    @PromptServer.instance.routes.post("/o1key/video/prompt-write")
    async def write_o1key_video_prompt(request):
        import aiohttp as _aiohttp

        try:
            data = await request.json()
            if not isinstance(data, dict):
                raise ValueError("请求体必须是对象")
            api_key = get_api_key() or ""
            if not api_key:
                return web.json_response({"error": "未配置 API Key"}, status=401)

            context = {
                "generation_mode": data.get("generation_mode", "text"),
                "duration": data.get("duration", "auto"),
                "aspect_ratio": data.get("aspect_ratio", "auto"),
                "generate_audio": data.get("generate_audio", False),
                "reference_video_count": data.get("reference_video_count", 0),
                "reference_audio_count": data.get("reference_audio_count", 0),
            }
            timeout = _aiohttp.ClientTimeout(total=PROMPT_OPTIMIZER_TIMEOUT_SECONDS + 15)
            async with _aiohttp.ClientSession(timeout=timeout) as session:
                written = await write_video_prompt(
                    session,
                    NETWORK_ROUTES[get_network_route()],
                    api_key,
                    data.get("prompt", ""),
                    data.get("references", []),
                    folder_paths.get_input_directory(),
                    context,
                )
            return web.json_response(
                {
                    "prompt": written,
                    "model": "gpt-5.6-sol",
                    "reasoning_effort": "high",
                    "preset": "video-default",
                },
                headers={"Cache-Control": "no-store"},
            )
        except ValueError as exc:
            return web.json_response({"error": str(exc)}, status=400)
        except RuntimeError as exc:
            return web.json_response({"error": str(exc)}, status=502)
        except Exception:
            return web.json_response({"error": "视频 AI帮写失败，请稍后重试"}, status=500)

    # === AI 聊天代理（流式 SSE 透传） ===
    @PromptServer.instance.routes.post("/o1key/chat/completions")
    async def chat_completions_proxy(request):
        import aiohttp as _aiohttp
        import json as _cjson

        data = await request.json()
        config = load_config()
        api_key = config.get("O1KEY_API_KEY", "")
        if not api_key:
            return web.json_response({"error": "未配置 API Key"}, status=401)

        base_url = NETWORK_ROUTES[get_network_route()]
        model = data.get("model", "gpt-6-sol")
        messages = data.get("messages", [])
        if not isinstance(messages, list) or not messages:
            return web.json_response({"error": "缺少对话内容"}, status=400)

        try:
            messages = expand_xlsx_attachments(messages)
        except ValueError as exc:
            return web.json_response({"error": str(exc)}, status=400)

        url = f"{base_url}/v1/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        }
        reasoning = data.get("reasoning_effort", "high")
        if reasoning not in ("low", "medium", "high"):
            reasoning = "high"
        body = {"model": model, "messages": messages, "stream": True}
        if model == "claude-fable-5":
            budgets = {"low": 2048, "medium": 8192, "high": 16384}
            budget = budgets[reasoning]
            body["thinking"] = {"type": "enabled", "budget_tokens": budget}
            body["max_tokens"] = budget + 8192
        elif model in ("gpt-5.5", "gpt-5.6-sol", "gpt-6-astra", "gpt-6-sol", "gemini-3.1-pro-preview"):
            body["reasoning_effort"] = reasoning

        search_trace = None
        timeout = _aiohttp.ClientTimeout(total=120)
        async with _aiohttp.ClientSession(timeout=timeout) as session:
            if data.get("web_search") is True:
                raw_query = extract_search_query(messages)
                if raw_query:
                    query = await rewrite_search_query(session, base_url, api_key, raw_query) or raw_query
                    try:
                        results = await web_search(session, query)
                        search_trace = {
                            "query": query,
                            "results": [
                                {"title": item["title"], "url": item["url"]}
                                for item in results
                            ],
                        }
                        messages = list(messages)
                        messages.insert(max(0, len(messages) - 1), {
                            "role": "system",
                            "content": build_search_context(query, results),
                        })
                        body["messages"] = messages
                    except Exception as exc:
                        search_trace = {"query": query, "results": [], "error": str(exc)}

            resp = web.StreamResponse(
                status=200, reason="OK",
                headers={
                    "Content-Type": "text/event-stream",
                    "Cache-Control": "no-cache",
                    "X-Accel-Buffering": "no",
                }
            )
            await resp.prepare(request)

            try:
                if search_trace:
                    event = _cjson.dumps({"o1key_search": search_trace}, ensure_ascii=False)
                    await resp.write(f"data: {event}\n\n".encode("utf-8"))
                async with session.post(url, headers=headers, json=body) as upstream:
                    if upstream.status != 200:
                        err = await upstream.text()
                        event = _cjson.dumps({"error": err}, ensure_ascii=False)
                        await resp.write(f"data: {event}\n\n".encode("utf-8"))
                        await resp.write(b"data: [DONE]\n\n")
                        return resp
                    async for chunk in upstream.content.iter_any():
                        await resp.write(chunk)
            except Exception as e:
                event = _cjson.dumps({"error": str(e)}, ensure_ascii=False)
                await resp.write(f"data: {event}\n\n".encode("utf-8"))
                await resp.write(b"data: [DONE]\n\n")

        return resp
    # === 执行事件 Hook：持久化耗时元数据 ===
    import time as _time, json as _json2, os as _os
    _execution_tracker = {}

    _orig_send_sync = PromptServer.instance.send_sync

    def _patched_send_sync(event, data, *args, **kwargs):
        try:
            if event == "execution_start":
                pid = data.get("prompt_id", "")
                if pid:
                    _execution_tracker[pid] = {"start": _time.time(), "outputs": []}
            elif event == "executed":
                pid = data.get("prompt_id", "")
                output = data.get("output") or {}
                if pid and pid in _execution_tracker:
                    for img in output.get("images", []) + output.get("gifs", []):
                        if img.get("type") == "output" and img.get("filename"):
                            _execution_tracker[pid]["outputs"].append(img["filename"])
            elif event == "executing" and data.get("node") is None:
                pid = data.get("prompt_id", "")
                tracker = _execution_tracker.pop(pid, None)
                if tracker and tracker["outputs"]:
                    end_time = _time.time()
                    start_time = tracker["start"]
                    output_dir = _os.path.abspath(folder_paths.get_output_directory())
                    meta_file = _get_o1key_history_meta_file(output_dir)
                    meta = {}
                    if _os.path.isfile(meta_file):
                        try:
                            with open(meta_file, "r", encoding="utf-8") as mf:
                                meta = _json2.load(mf)
                        except Exception:
                            pass
                    for fname in tracker["outputs"]:
                        meta[fname] = {
                            "start_time": start_time,
                            "end_time": end_time,
                            "outputs_count": len(tracker["outputs"]),
                            "workflow_id": pid,
                        }
                    try:
                        with open(meta_file, "w", encoding="utf-8") as mf:
                            _json2.dump(meta, mf, ensure_ascii=False)
                    except Exception:
                        pass
        except Exception:
            pass
        return _orig_send_sync(event, data, *args, **kwargs)

    PromptServer.instance.send_sync = _patched_send_sync

except Exception:
    pass