ComfyUI-OpenClaw/__init__.py

import os
import sys

# Ensure this custom node root is on sys.path (ComfyUI loads modules by path, not package)
_MOLTBOT_ROOT = os.path.dirname(os.path.abspath(__file__))
if _MOLTBOT_ROOT not in sys.path:
    sys.path.insert(0, _MOLTBOT_ROOT)

if __package__:
    from .nodes.batch_variants import MoltbotBatchVariants
    from .nodes.image_to_prompt import MoltbotImageToPrompt
    from .nodes.prompt_planner import MoltbotPromptPlanner
    from .nodes.prompt_refiner import MoltbotPromptRefiner

    NODE_CLASS_MAPPINGS = {
        "MoltbotPromptPlanner": MoltbotPromptPlanner,
        "MoltbotBatchVariants": MoltbotBatchVariants,
        "MoltbotImageToPrompt": MoltbotImageToPrompt,
        "MoltbotPromptRefiner": MoltbotPromptRefiner,
    }

    NODE_DISPLAY_NAME_MAPPINGS = {
        "MoltbotPromptPlanner": "openclaw: Prompt Planner",
        "MoltbotBatchVariants": "openclaw: Batch Variants",
        "MoltbotImageToPrompt": "openclaw: Image to Prompt",
        "MoltbotPromptRefiner": "openclaw: Prompt Refiner",
    }

    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
else:
    # Allow test collection to proceed without crashing on relative imports
    NODE_CLASS_MAPPINGS = {}
    NODE_DISPLAY_NAME_MAPPINGS = {}
    __all__ = ["WEB_DIRECTORY"]

WEB_DIRECTORY = "./web"

# Register API routes (observability endpoints)
_routes_registered = False

import logging
import threading
import time


def _register_routes_once():
    global _routes_registered
    if _routes_registered:
        return

    # R23: Register Plugins
    # This needs to happen regardless of PromptServer availability, as plugins might register other things.
    try:
        from .services.plugins.builtin import register_all

        register_all()
    except Exception as e:
        logging.getLogger("ComfyUI-OpenClaw").error(f"Failed to register plugins: {e}")

    # R63/R84: Initialize Service & Module Registries
    try:
        from .services.modules import ModuleCapability, ModuleRegistry, enable_module
        from .services.registry import SVC_RUNTIME_CONFIG, ServiceRegistry
        from .services.runtime_config import get_config

        # 1. Register Runtime Config
        config = get_config()
        ServiceRegistry.register(SVC_RUNTIME_CONFIG, config)

        # 2. Initialize Default Modules
        # Always enable CORE and SECURITY
        enable_module(ModuleCapability.CORE)
        enable_module(ModuleCapability.SECURITY)
        enable_module(ModuleCapability.OBSERVABILITY)

        # 3. Conditional Modules (Config-driven)
        # Bridge
        if config.bridge_enabled:
            enable_module(ModuleCapability.BRIDGE)

        # Scheduler (Always enabled as core feature, but runner starts conditionally)
        enable_module(ModuleCapability.SCHEDULER)

        # Webhook (Always enabled as core feature, but auth-gated)
        enable_module(ModuleCapability.WEBHOOK)

        # Connector (Always enabled as core feature)
        enable_module(ModuleCapability.CONNECTOR)

        # Lock registry alignment
        ModuleRegistry.lock()
        logging.getLogger("ComfyUI-OpenClaw").info(
            f"Initialized modules: {ModuleRegistry.get_enabled_list()}"
        )

        # R41/S41: Run Security Gate
        # Must run after config and modules are loaded.
        from .services.security_gate import enforce_startup_gate

        enforce_startup_gate()

    except Exception as e:
        logging.getLogger("ComfyUI-OpenClaw").error(
            f"Failed to initialize registries: {e}"
        )

    def _do_full_registration(server):
        """Register all Moltbot routes including Bridge and Scheduler."""
        from .api.approvals import register_approval_routes
        from .api.bridge import BridgeHandlers
        from .api.presets import register_preset_routes
        from .api.routes import register_routes
        from .api.schedules import register_schedule_routes
        from .api.triggers import register_trigger_routes
        from .services.access_control import require_admin_token
        from .services.plugins.async_bridge import run_async_in_sync_context
        from .services.queue_submit import submit_prompt
        from .services.templates import get_template_service

        register_routes(server)
        register_preset_routes(server.app)

        # R4: Register schedule CRUD routes
        register_schedule_routes(server.app, require_admin_token_fn=require_admin_token)

        # Bridge: Adapt functional submit_prompt to service interface
        class QueueSubmitService:
            def submit(self, job_req):
                tmpl_svc = get_template_service()
                workflow = tmpl_svc.render_template(job_req.template_id, job_req.inputs)

                async def _do_submit():
                    return await submit_prompt(
                        workflow,
                        client_id=job_req.session_id or "bridge",
                        extra_data={
                            "openclaw": {"trace_id": job_req.trace_id},
                            # Legacy key kept for existing tooling that expects this blob.
                            "moltbot": {"trace_id": job_req.trace_id},
                        },
                        source="bridge",
                        trace_id=job_req.trace_id,
                    )

                return run_async_in_sync_context(_do_submit())

        bridge_handlers = BridgeHandlers(submit_service=QueueSubmitService())
        # Dual registration for bridge (R26)
        if hasattr(server.app.router, "add_post"):  # Redundant check but safe
            # Legacy
            server.app.router.add_post(
                "/moltbot/bridge/submit", bridge_handlers.submit_handler
            )
            server.app.router.add_post(
                "/moltbot/bridge/deliver", bridge_handlers.deliver_handler
            )
            server.app.router.add_get(
                "/moltbot/bridge/health", bridge_handlers.health_handler
            )
            # New prefix
            server.app.router.add_post(
                "/openclaw/bridge/submit", bridge_handlers.submit_handler
            )
            server.app.router.add_post(
                "/openclaw/bridge/deliver", bridge_handlers.deliver_handler
            )
            server.app.router.add_get(
                "/openclaw/bridge/health", bridge_handlers.health_handler
            )
            # /api Prefixed
            try:
                server.app.router.add_post(
                    "/api/moltbot/bridge/submit", bridge_handlers.submit_handler
                )
                server.app.router.add_post(
                    "/api/moltbot/bridge/deliver", bridge_handlers.deliver_handler
                )
                server.app.router.add_get(
                    "/api/moltbot/bridge/health", bridge_handlers.health_handler
                )
                server.app.router.add_post(
                    "/api/openclaw/bridge/submit", bridge_handlers.submit_handler
                )
                server.app.router.add_post(
                    "/api/openclaw/bridge/deliver", bridge_handlers.deliver_handler
                )
                server.app.router.add_get(
                    "/api/openclaw/bridge/health", bridge_handlers.health_handler
                )
            except RuntimeError:
                pass

        # Shared submit function for scheduler and triggers
        async def unified_submit_fn(
            template_id,
            inputs,
            trace_id,
            idempotency_key,
            delivery=None,
            source="unknown",
        ):
            """Submit function for scheduler and trigger-triggered runs."""
            # NOTE: Use IdempotencyStore API (check_and_record/update_prompt_id).
            # Avoid legacy get_store/get/set usage; wrong API here breaks route registration at runtime.
            from .services.idempotency_store import IdempotencyStore
            from .services.queue_submit import submit_prompt
            from .services.templates import get_template_service

            # Check idempotency
            store = IdempotencyStore()
            is_dup, existing_prompt_id = store.check_and_record(idempotency_key)
            if is_dup:
                return {"prompt_id": existing_prompt_id, "deduped": True}

            # Render template
            tmpl_svc = get_template_service()
            workflow = tmpl_svc.render_template(template_id, inputs)

            # Submit
            result = await submit_prompt(
                workflow,
                extra_data={
                    "openclaw": {"trace_id": trace_id, "source": "automation"},
                    "moltbot": {"trace_id": trace_id, "source": "automation"},
                },
                source=source,
                trace_id=trace_id,
            )

            # Store for dedupe
            if result.get("prompt_id"):
                store.update_prompt_id(idempotency_key, result["prompt_id"])

            return result

        # R4: Start scheduler daemon
        from .services.scheduler.runner import get_scheduler_runner, start_scheduler

        runner = get_scheduler_runner()
        runner._submit_fn = unified_submit_fn
        start_scheduler()

        # F6: Register trigger routes
        register_trigger_routes(
            server.app,
            require_admin_token_fn=require_admin_token,
            submit_fn=unified_submit_fn,
        )

        # S7: Register approval routes (with execution capability)
        register_approval_routes(
            server.app,
            require_admin_token_fn=require_admin_token,
            submit_fn=unified_submit_fn,
        )

    def start_registration_retry_loop():
        """
        R25: Background loop to ensure routes are registered even if PromptServer is slow to init.
        Attempts 10 times with exponential backoff.
        """

        def _retry_worker():
            global _routes_registered
            attempts = 0
            max_attempts = 10
            delay = 2.0
            logger = logging.getLogger("ComfyUI-OpenClaw")

            while not _routes_registered and attempts < max_attempts:
                try:
                    ps_mod = sys.modules.get("server")
                    PromptServer = (
                        getattr(ps_mod, "PromptServer", None) if ps_mod else None
                    )
                    if (
                        PromptServer
                        and getattr(PromptServer, "instance", None) is not None
                    ):
                        _do_full_registration(PromptServer.instance)
                        _routes_registered = True
                        logger.info(
                            f"Routes registered successfully on attempt {attempts + 1}"
                        )
                        return
                    logger.debug(
                        f"PromptServer.instance not ready (attempt {attempts + 1})"
                    )
                except Exception as e:
                    logger.exception(
                        f"Error registering routes (attempt {attempts + 1})"
                    )

                time.sleep(delay)
                delay = min(delay * 1.5, 30)
                attempts += 1

            if not _routes_registered:
                logger.error(
                    f"Failed to register routes after {max_attempts} attempts. API endpoints unavailable."
                )

        t = threading.Thread(
            target=_retry_worker, name="openclaw-route-retry", daemon=True
        )
        t.start()

    # Initial attempt + Start background retry
    try:
        ps_mod = sys.modules.get("server")
        PromptServer = getattr(ps_mod, "PromptServer", None) if ps_mod else None

        if PromptServer and getattr(PromptServer, "instance", None) is not None:
            _do_full_registration(PromptServer.instance)
            _routes_registered = True
            logging.getLogger("ComfyUI-OpenClaw").info(
                "Routes registered successfully on initial attempt."
            )
        else:
            logging.getLogger("ComfyUI-OpenClaw").info(
                "PromptServer not ready, starting background registration retry loop..."
            )
            start_registration_retry_loop()

    except Exception as e:
        logging.getLogger("ComfyUI-OpenClaw").exception("Route registration failed")


_register_routes_once()