ComfyUI-OpenClaw/api/assist.py

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
from typing import Any, Dict, Optional

try:
    from ..services.access_control import require_admin_token
    from ..services.aiohttp_compat import import_aiohttp_web
    from ..services.async_utils import run_in_thread
    from ..services.automation_composer import AutomationComposerService
    from ..services.planner import PlannerService
    from ..services.planner_registry import get_planner_registry
    from ..services.rate_limit import build_rate_limit_response, check_rate_limit
    from ..services.reasoning_redaction import (
        audit_reasoning_reveal,
        resolve_reasoning_reveal,
        sanitize_operator_payload,
    )
    from ..services.refiner import RefinerService
except ImportError:
    # Fallback for ComfyUI's non-package loader or ad-hoc imports.
    from services.access_control import require_admin_token
    from services.aiohttp_compat import import_aiohttp_web
    from services.async_utils import run_in_thread
    from services.automation_composer import AutomationComposerService
    from services.planner import PlannerService
    from services.planner_registry import get_planner_registry
    from services.rate_limit import build_rate_limit_response, check_rate_limit
    from services.reasoning_redaction import (
        audit_reasoning_reveal,
        resolve_reasoning_reveal,
        sanitize_operator_payload,
    )
    from services.refiner import RefinerService

# R98: Endpoint Metadata
if __package__ and "." in __package__:
    from ..services.endpoint_manifest import (
        AuthTier,
        RiskTier,
        RoutePlane,
        endpoint_metadata,
    )
else:
    from services.endpoint_manifest import (
        AuthTier,
        RiskTier,
        RoutePlane,
        endpoint_metadata,
    )

logger = logging.getLogger("ComfyUI-OpenClaw.api.assist")
web = import_aiohttp_web()

# Payload size limits (character count for strings, base64 length for images)
MAX_REQUIREMENTS_LEN = 8000
MAX_STYLE_LEN = 2000
MAX_IMAGE_B64_LEN = 5 * 1024 * 1024  # ~5MB base64 string length
MAX_STREAM_DELTA_CHARS = 256
MAX_STREAM_PREVIEW_CHARS = 16_000
STREAM_KEEPALIVE_SEC = 1.0


def _planner_profiles_payload() -> Dict[str, Any]:
    registry = get_planner_registry()
    return {
        "profiles": [
            {
                "id": profile.id,
                "label": profile.label,
                "description": profile.description,
                "version": profile.version,
            }
            for profile in registry.list_profiles()
        ],
        "default_profile": registry.get_default_profile_id(),
    }


class AssistHandlers:
    def __init__(self):
        self.planner = PlannerService()
        self.refiner = RefinerService()
        self.composer = AutomationComposerService()

    async def _require_admin_and_rate_limit(
        self, request: web.Request
    ) -> Optional[web.Response]:
        authorized, _err_msg = require_admin_token(request)
        if not authorized:
            return web.json_response({"error": "Unauthorized"}, status=401)
        if not check_rate_limit(request, "admin"):
            return build_rate_limit_response(
                request,
                "admin",
                web_module=web,
                error="Rate limit exceeded",
                include_ok=False,
            )
        return None

    async def _parse_json_body(
        self, request: web.Request
    ) -> tuple[Optional[dict], Optional[web.Response]]:
        try:
            data = await request.json()
        except Exception:
            return None, web.json_response({"error": "Invalid JSON"}, status=400)
        if not isinstance(data, dict):
            return None, web.json_response(
                {"error": "JSON object required"}, status=400
            )
        return data, None

    def _validate_planner_payload(
        self, data: dict
    ) -> tuple[Optional[dict], Optional[web.Response]]:
        registry = get_planner_registry()
        profile = data.get("profile", registry.get_default_profile_id())
        requirements = data.get("requirements", "")
        style = data.get("style_directives", "")
        seed = data.get("seed", 0)

        if not isinstance(profile, str):
            return None, web.json_response(
                {"error": "profile must be string"}, status=400
            )
        if not registry.get_profile(profile):
            return None, web.json_response(
                {"error": f"Unknown profile: {profile}"}, status=400
            )
        if not isinstance(requirements, str):
            return None, web.json_response(
                {"error": "requirements must be string"}, status=400
            )
        if not isinstance(style, str):
            return None, web.json_response(
                {"error": "style_directives must be string"}, status=400
            )
        if len(requirements) > MAX_REQUIREMENTS_LEN:
            return None, web.json_response(
                {"error": f"requirements exceeds {MAX_REQUIREMENTS_LEN} chars"},
                status=400,
            )
        if len(style) > MAX_STYLE_LEN:
            return None, web.json_response(
                {"error": f"style_directives exceeds {MAX_STYLE_LEN} chars"}, status=400
            )
        try:
            seed = int(seed)
        except Exception:
            seed = 0
        return {
            "profile": profile,
            "requirements": requirements,
            "style_directives": style,
            "seed": seed,
        }, None

    def _validate_refiner_payload(
        self, data: dict
    ) -> tuple[Optional[dict], Optional[web.Response]]:
        image_b64 = data.get("image_b64", "")
        orig_pos = data.get("orig_positive", "")
        orig_neg = data.get("orig_negative", "")
        issue = data.get("issue", "Fix issues")
        params_json = data.get("params_json", "{}")
        goal = data.get("goal", "Fix issues")

        if not isinstance(image_b64, str) or not image_b64:
            return None, web.json_response({"error": "image_b64 required"}, status=400)
        if len(image_b64) > MAX_IMAGE_B64_LEN:
            return None, web.json_response(
                {"error": f"image_b64 exceeds {MAX_IMAGE_B64_LEN // 1024 // 1024}MB"},
                status=400,
            )
        for key, value in (
            ("orig_positive", orig_pos),
            ("orig_negative", orig_neg),
            ("issue", issue),
            ("params_json", params_json),
            ("goal", goal),
        ):
            if not isinstance(value, str):
                return None, web.json_response(
                    {"error": f"{key} must be string"}, status=400
                )
        if len(orig_pos) > MAX_REQUIREMENTS_LEN or len(orig_neg) > MAX_REQUIREMENTS_LEN:
            return None, web.json_response({"error": "Prompt too long"}, status=400)

        return {
            "image_b64": image_b64,
            "orig_positive": orig_pos,
            "orig_negative": orig_neg,
            "issue": issue,
            "params_json": params_json,
            "goal": goal,
        }, None

    def _finalize_operator_payload(
        self,
        request: web.Request,
        *,
        target: str,
        payload: Dict[str, Any],
        service: Any,
    ) -> Dict[str, Any]:
        admin_allowed, _ = require_admin_token(request)
        reveal = resolve_reasoning_reveal(request, admin_authorized=admin_allowed)
        audit_reasoning_reveal(request, target=target, decision=reveal)

        final_payload = sanitize_operator_payload(payload)
        consume_debug = getattr(service, "consume_last_reasoning_debug", None)
        reasoning_debug = consume_debug() if callable(consume_debug) else None
        if reveal["allowed"] and reasoning_debug not in (None, {}, []):
            final_payload = dict(final_payload)
            final_payload["debug"] = {"reasoning": reasoning_debug}
        return final_payload

    @staticmethod
    def _sse_frame(event: str, payload: Dict[str, Any]) -> bytes:
        return (
            f"event: {event}\n"
            f"data: {json.dumps(payload, ensure_ascii=False, separators=(',', ':'))}\n\n"
        ).encode("utf-8")

    async def _write_sse_event(
        self, response: web.StreamResponse, event: str, payload: Dict[str, Any]
    ) -> bool:
        try:
            await response.write(self._sse_frame(event, payload))
            return True
        except (ConnectionError, RuntimeError):
            return False

    async def _assist_stream_session(
        self,
        request: web.Request,
        *,
        kind: str,
        worker_fn,
        worker_kwargs: Dict[str, Any],
    ) -> web.StreamResponse:
        response = web.StreamResponse(
            status=200,
            headers={
                "Content-Type": "text/event-stream",
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )
        await response.prepare(request)

        loop = asyncio.get_running_loop()
        queue: asyncio.Queue = asyncio.Queue()
        preview_chars = 0

        def emit(event: str, payload: Dict[str, Any]) -> None:
            try:
                loop.call_soon_threadsafe(queue.put_nowait, (event, payload))
            except RuntimeError:
                pass

        def on_text_delta(delta: str) -> None:
            nonlocal preview_chars
            if not isinstance(delta, str) or not delta:
                return
            remaining = MAX_STREAM_PREVIEW_CHARS - preview_chars
            if remaining <= 0:
                return
            clipped = delta[: min(remaining, MAX_STREAM_DELTA_CHARS)]
            if not clipped:
                return
            preview_chars += len(clipped)
            emit("delta", {"text": clipped, "preview_chars": preview_chars})

        async def runner() -> None:
            emit("ready", {"ok": True, "kind": kind, "mode": "sse"})
            emit(
                "stage", {"phase": "dispatch", "message": "Dispatching assist request"}
            )
            try:
                call_kwargs = dict(worker_kwargs)
                call_kwargs["on_text_delta"] = on_text_delta
                result = await run_in_thread(worker_fn, **call_kwargs)
                emit(
                    "stage",
                    {"phase": "finalize", "message": "Parsing and validating output"},
                )
                if kind == "planner":
                    pos, neg, params = result
                    final_payload = {
                        "positive": pos,
                        "negative": neg,
                        "params": params,
                    }
                elif kind == "refiner":
                    new_pos, new_neg, patch, rationale = result
                    final_payload = {
                        "refined_positive": new_pos,
                        "refined_negative": new_neg,
                        "param_patch": patch,
                        "rationale": rationale,
                    }
                else:
                    final_payload = {"result": result}
                service = (
                    self.planner
                    if kind == "planner"
                    else self.refiner if kind == "refiner" else self.composer
                )
                final_payload = self._finalize_operator_payload(
                    request,
                    target=f"assist.{kind}.stream",
                    payload=final_payload,
                    service=service,
                )
                emit(
                    "final",
                    {
                        "ok": True,
                        "kind": kind,
                        "result": final_payload,
                        "streaming": {
                            "preview_chars": preview_chars,
                            "preview_truncated": preview_chars
                            >= MAX_STREAM_PREVIEW_CHARS,
                        },
                    },
                )
            except Exception as e:
                logger.exception("Assist streaming API failed (%s)", kind)
                emit(
                    "error",
                    {"ok": False, "kind": kind, "error": "Internal server error"},
                )
            finally:
                emit("__done__", {})

        runner_task = asyncio.create_task(runner())

        try:
            while True:
                try:
                    event, payload = await asyncio.wait_for(
                        queue.get(), timeout=STREAM_KEEPALIVE_SEC
                    )
                except asyncio.TimeoutError:
                    if runner_task.done():
                        break
                    if not await self._write_sse_event(
                        response, "keepalive", {"ok": True}
                    ):
                        break
                    continue

                if event == "__done__":
                    break
                if not await self._write_sse_event(response, event, payload):
                    break
        finally:
            if not runner_task.done():
                runner_task.cancel()
                with contextlib.suppress(BaseException):
                    await runner_task
        return response

    @endpoint_metadata(
        auth=AuthTier.ADMIN,
        risk=RiskTier.LOW,
        summary="List planner profiles",
        description="Returns Prompt Planner profiles from the active registry.",
        audit="assist.planner_profiles",
        plane=RoutePlane.ADMIN,
    )
    async def planner_profiles_handler(self, request):
        auth_resp = await self._require_admin_and_rate_limit(request)
        if auth_resp:
            return auth_resp
        return web.json_response(_planner_profiles_payload())

    @endpoint_metadata(
        auth=AuthTier.ADMIN,
        risk=RiskTier.MEDIUM,
        summary="Run planner",
        description="Generate prompts from requirements via LLM.",
        audit="assist.planner",
        plane=RoutePlane.ADMIN,
    )
    async def planner_handler(self, request):
        """
        POST /openclaw/assist/planner (legacy: /moltbot/assist/planner)
        JSON: { profile, requirements, style_directives, seed }
        """
        # Security: Admin Token required
        auth_resp = await self._require_admin_and_rate_limit(request)
        if auth_resp:
            return auth_resp
        data, error_resp = await self._parse_json_body(request)
        if error_resp:
            return error_resp
        assert data is not None
        payload, payload_err = self._validate_planner_payload(data)
        if payload_err:
            return payload_err
        assert payload is not None

        try:
            # Run sync LLM call in thread pool to avoid blocking event loop
            pos, neg, params = await run_in_thread(
                self.planner.plan_generation,
                payload["profile"],
                payload["requirements"],
                payload["style_directives"],
                payload["seed"],
            )

            return web.json_response(
                self._finalize_operator_payload(
                    request,
                    target="assist.planner",
                    payload={"positive": pos, "negative": neg, "params": params},
                    service=self.planner,
                )
            )

        except Exception as e:
            logger.exception("Planner API failed")
            return web.json_response({"error": "Internal server error"}, status=500)

    @endpoint_metadata(
        auth=AuthTier.ADMIN,
        risk=RiskTier.MEDIUM,
        summary="Run refiner",
        description="Refine prompt/parameters based on feedback.",
        audit="assist.refiner",
        plane=RoutePlane.ADMIN,
    )
    async def refiner_handler(self, request):
        """
        POST /openclaw/assist/refiner (legacy: /moltbot/assist/refiner)
        JSON: { image_b64, orig_positive, orig_negative, issue, params_json, goal }
        """
        # Security checks
        auth_resp = await self._require_admin_and_rate_limit(request)
        if auth_resp:
            return auth_resp
        data, error_resp = await self._parse_json_body(request)
        if error_resp:
            return error_resp
        assert data is not None
        payload, payload_err = self._validate_refiner_payload(data)
        if payload_err:
            return payload_err
        assert payload is not None

        try:
            # Run sync LLM call in thread pool
            new_pos, new_neg, patch, rationale = await run_in_thread(
                self.refiner.refine_prompt,
                **payload,
            )

            return web.json_response(
                self._finalize_operator_payload(
                    request,
                    target="assist.refiner",
                    payload={
                        "refined_positive": new_pos,
                        "refined_negative": new_neg,
                        "param_patch": patch,
                        "rationale": rationale,
                    },
                    service=self.refiner,
                )
            )
        except Exception as e:
            logger.exception("Refiner API failed")
            return web.json_response({"error": "Internal server error"}, status=500)

    @endpoint_metadata(
        auth=AuthTier.ADMIN,
        risk=RiskTier.MEDIUM,
        summary="Run planner (streaming)",
        description="Generate prompts from requirements via LLM with SSE-style incremental updates.",
        audit="assist.planner.stream",
        plane=RoutePlane.ADMIN,
    )
    async def planner_stream_handler(self, request):
        auth_resp = await self._require_admin_and_rate_limit(request)
        if auth_resp:
            return auth_resp
        data, error_resp = await self._parse_json_body(request)
        if error_resp:
            return error_resp
        assert data is not None
        payload, payload_err = self._validate_planner_payload(data)
        if payload_err:
            return payload_err
        assert payload is not None

        return await self._assist_stream_session(
            request,
            kind="planner",
            worker_fn=self.planner.plan_generation,
            worker_kwargs=payload,
        )

    @endpoint_metadata(
        auth=AuthTier.ADMIN,
        risk=RiskTier.MEDIUM,
        summary="Run refiner (streaming)",
        description="Refine prompt/parameters with SSE-style incremental updates.",
        audit="assist.refiner.stream",
        plane=RoutePlane.ADMIN,
    )
    async def refiner_stream_handler(self, request):
        auth_resp = await self._require_admin_and_rate_limit(request)
        if auth_resp:
            return auth_resp
        data, error_resp = await self._parse_json_body(request)
        if error_resp:
            return error_resp
        assert data is not None
        payload, payload_err = self._validate_refiner_payload(data)
        if payload_err:
            return payload_err
        assert payload is not None

        return await self._assist_stream_session(
            request,
            kind="refiner",
            worker_fn=self.refiner.refine_prompt,
            worker_kwargs=payload,
        )

    @endpoint_metadata(
        auth=AuthTier.ADMIN,
        risk=RiskTier.MEDIUM,
        summary="Compose automation payload",
        description="Generate-only automation payload draft for trigger/webhook endpoints.",
        audit="assist.compose",
        plane=RoutePlane.ADMIN,
    )
    async def compose_handler(self, request):
        """
        POST /openclaw/assist/automation/compose (legacy: /moltbot/assist/automation/compose)
        JSON:
        {
          kind: "trigger" | "webhook",
          template_id: str,
          intent: str,
          inputs_hint?: object,
          profile_id?: str,
          require_approval?: bool,
          trace_id?: str,
          callback?: object
        }
        """
        authorized, err_msg = require_admin_token(request)
        if not authorized:
            return web.json_response({"error": "Unauthorized"}, status=401)

        if not check_rate_limit(request, "admin"):
            return build_rate_limit_response(
                request,
                "admin",
                web_module=web,
                error="Rate limit exceeded",
                include_ok=False,
            )

        try:
            data = await request.json()
        except Exception:
            return web.json_response({"error": "Invalid JSON"}, status=400)

        kind = data.get("kind")
        template_id = data.get("template_id")
        intent = data.get("intent")
        inputs_hint = data.get("inputs_hint", {})
        profile_id = data.get("profile_id")
        require_approval = data.get("require_approval")
        trace_id = data.get("trace_id")
        callback = data.get("callback")

        if not isinstance(kind, str) or kind.strip().lower() not in {
            "trigger",
            "webhook",
        }:
            return web.json_response(
                {"error": "kind must be 'trigger' or 'webhook'"}, status=400
            )
        if not isinstance(template_id, str) or not template_id.strip():
            return web.json_response({"error": "template_id is required"}, status=400)
        if not isinstance(intent, str) or not intent.strip():
            return web.json_response({"error": "intent is required"}, status=400)
        if len(intent) > MAX_REQUIREMENTS_LEN:
            return web.json_response(
                {"error": f"intent exceeds {MAX_REQUIREMENTS_LEN} chars"}, status=400
            )
        if not isinstance(inputs_hint, dict):
            return web.json_response(
                {"error": "inputs_hint must be object"}, status=400
            )
        if profile_id is not None and not isinstance(profile_id, str):
            return web.json_response({"error": "profile_id must be string"}, status=400)
        if require_approval is not None and not isinstance(require_approval, bool):
            return web.json_response(
                {"error": "require_approval must be boolean"}, status=400
            )
        if trace_id is not None and not isinstance(trace_id, str):
            return web.json_response({"error": "trace_id must be string"}, status=400)
        if callback is not None and not isinstance(callback, dict):
            return web.json_response({"error": "callback must be object"}, status=400)

        try:
            result = await run_in_thread(
                self.composer.compose_payload,
                kind=kind,
                template_id=template_id,
                intent=intent,
                inputs_hint=inputs_hint,
                profile_id=profile_id,
                require_approval=require_approval,
                trace_id=trace_id,
                callback=callback,
            )
            return web.json_response(
                self._finalize_operator_payload(
                    request,
                    target="assist.compose",
                    payload={"ok": True, **result},
                    service=self.composer,
                )
            )
        except ValueError as e:
            return web.json_response({"error": str(e)}, status=400)
        except Exception:
            logger.exception("Automation compose API failed")
            return web.json_response({"error": "Internal server error"}, status=500)