research/widget_status.py

"""Capability-authorized HTTP snapshots for MCP App browser polling."""

from __future__ import annotations

from urllib.parse import quote

from fastmcp import FastMCP
from fastmcp.apps.app import _make_resolver
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from .app_jobs import ResearchJobStore
from .app_ui import _ui_snapshot, build_research_ui
from .hf_design import HFDesign
from .status_capability import StatusCapabilityStore

APP_NAME = "Hugging Face Researcher"
STATUS_PATH = "/widget/research/{token}"
RESPONSE_HEADERS = {
    "Access-Control-Allow-Origin": "*",
    "Cache-Control": "no-store",
    "Referrer-Policy": "no-referrer",
    "X-Content-Type-Options": "nosniff",
}


def capability_urls(origin: str, token: str) -> tuple[str, str]:
    status_url = f"{origin.rstrip('/')}/widget/research/{quote(token, safe='')}"
    return status_url, f"{status_url}?view=1"


def register_widget_status_route(
    mcp: FastMCP,
    *,
    jobs: ResearchJobStore,
    capabilities: StatusCapabilityStore,
    build_id: str,
    origin: str,
    design: HFDesign,
) -> None:
    @mcp.custom_route(STATUS_PATH, methods=["GET", "OPTIONS"])
    async def widget_status(request: Request) -> Response:
        if request.method == "OPTIONS":
            return Response(
                status_code=204,
                headers={
                    **RESPONSE_HEADERS,
                    "Access-Control-Allow-Methods": "GET, OPTIONS",
                },
            )

        capability = capabilities.resolve(request.path_params["token"])
        if capability is None:
            return JSONResponse(
                {"detail": "Unavailable"},
                status_code=404,
                headers=RESPONSE_HEADERS,
            )
        job = await jobs.get(capability.job_id, capability.owner_id)
        if job is None:
            return JSONResponse(
                {"detail": "Unavailable"},
                status_code=404,
                headers=RESPONSE_HEADERS,
            )

        if request.query_params.get("view") == "1":
            status_url, recovery_url = capability_urls(
                origin, request.path_params["token"]
            )
            snapshot = job.snapshot()
            snapshot["markdown_report_blocks"] = job.markdown_report_blocks
            app = build_research_ui(
                job.topic,
                snapshot,
                build_id=build_id,
                design=design,
                status_url=status_url,
                recovery_url=recovery_url,
            )
            payload = app.to_json(tool_resolver=_make_resolver(APP_NAME))
        else:
            payload = _ui_snapshot(job.snapshot())
            payload["report_blocks"] = job.markdown_report_blocks
            payload["report_ready"] = bool(job.markdown_report_blocks)

        return JSONResponse(payload, headers=RESPONSE_HEADERS)