ComfyAssets-ComfyUI_PromptManager/py/api/job_routes.py

"""Session job routes (Activity): list, inspect, follow and cancel jobs.

The registry itself is ``py/api/jobs.py``. The scan job's runner is in
``py/api/admin.py``, the image tagging job's in ``py/api/image_tag_jobs.py``.
Every response leaves out absolute server paths: job rows already carry
root-relative paths and public root labels.
"""

import contextlib
import json

from aiohttp import web

from .jobs import DEFAULT_SCAN_SCOPE, SCAN_SCOPES, JobBusy


def _not_found():
    return web.json_response({"success": False, "error": "Job not found"}, status=404)


def _bad_request(message):
    return web.json_response({"success": False, "error": message}, status=400)


async def _scan_scope(request):
    """The requested scan scope ("new" or "all"), or None when the body is invalid.

    An empty body (also a chunked one with no bytes), or one without
    ``scope``, asks for the default.
    """
    if not request.can_read_body:
        return DEFAULT_SCAN_SCOPE
    text = await request.text()
    if not text.strip():  # e.g. a chunked POST that sent no bytes
        return DEFAULT_SCAN_SCOPE
    try:
        body = json.loads(text)
    except ValueError:
        return None
    if not isinstance(body, dict):
        return None
    scope = body.get("scope", DEFAULT_SCAN_SCOPE)
    return scope if scope in SCAN_SCOPES else None


def busy_response(busy: JobBusy):
    """``409 busy``, naming the running job that holds the slot."""
    running = busy.running
    return web.json_response(
        {
            "success": False,
            "code": "busy",
            "error": str(busy),
            "running": {"id": running.id, "title": running.title},
        },
        status=409,
    )


def _summary_event(job):
    return f"data: {json.dumps(job.summary())}\n\n".encode("utf-8")


class JobRoutesMixin:
    """Mixin providing the /prompt_manager/jobs endpoints."""

    def _register_job_routes(self, routes):
        @routes.get("/prompt_manager/jobs")
        async def list_jobs_route(request):
            return await self.list_jobs(request)

        @routes.post("/prompt_manager/jobs/scan")
        async def start_scan_job_route(request):
            return await self.start_scan_job(request)

        @routes.post("/prompt_manager/jobs/image-tags")
        async def start_image_tag_job_route(request):
            return await self.start_image_tag_job(request)

        @routes.get("/prompt_manager/jobs/{job_id}")
        async def job_detail_route(request):
            return await self.job_detail(request)

        @routes.get("/prompt_manager/jobs/{job_id}/report")
        async def job_report_route(request):
            return await self.job_report(request)

        @routes.get("/prompt_manager/jobs/{job_id}/events")
        async def job_events_route(request):
            return await self.job_events(request)

        @routes.post("/prompt_manager/jobs/{job_id}/cancel")
        async def cancel_job_route(request):
            return await self.cancel_job(request)

    async def list_jobs(self, request):
        """This session's jobs: running first, then finished, newest first."""
        jobs = [job.summary() for job in self._jobs.list()]
        return web.json_response({"success": True, "jobs": jobs})

    async def job_detail(self, request):
        job = self._jobs.get(request.match_info["job_id"])
        if job is None:
            return _not_found()
        return web.json_response({"success": True, "job": job.detail()})

    async def job_report(self, request):
        """The job's whole record as a JSON attachment (no absolute paths)."""
        job = self._jobs.get(request.match_info["job_id"])
        if job is None:
            return _not_found()
        body = json.dumps(job.report(), indent=2)
        disposition = f'attachment; filename="{job.report_filename()}"'
        return web.Response(
            text=body,
            content_type="application/json",
            headers={
                "Content-Disposition": disposition,
                "Cache-Control": "no-store",
            },
        )

    async def cancel_job(self, request):
        """Stop a job at its next file boundary; a finished job is a no-op."""
        job_id = request.match_info["job_id"]
        if self._jobs.get(job_id) is None:
            return _not_found()
        requested = self._jobs.cancel(job_id)
        return web.json_response({"success": True, "cancellation_requested": requested})

    async def start_scan_job(self, request):
        """Start the metadata scan as a job (202), or 409 naming the busy job.

        The body's ``scope`` is "new" (the default: new and unresolved files)
        or "all"; anything else is a 400.
        """
        files = await _scan_scope(request)
        if files is None:
            return _bad_request('The scan scope must be "new" or "all"')
        try:
            job = self._start_scan_job(files)
        except JobBusy as busy:
            return busy_response(busy)
        return web.json_response({"success": True, "job": job.summary()}, status=202)

    async def job_events(self, request):
        """Stream the job's summary as SSE: now, after each event, and at the end."""
        job = self._jobs.get(request.match_info["job_id"])
        if job is None:
            return _not_found()
        response = web.StreamResponse(
            status=200,
            reason="OK",
            headers={
                "Content-Type": "text/event-stream",
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            },
        )
        await response.prepare(request)
        try:
            await response.write(_summary_event(job))
            async with contextlib.aclosing(job.events()) as events:
                async for _payload in events:
                    await response.write(_summary_event(job))
            # The last event is published before the job takes its final state.
            await response.write(_summary_event(job))
            await response.write_eof()
        except Exception as e:  # the socket closed under us; the job goes on
            self.logger.info(f"Job progress client disconnected ({e})")
        return response