trackio/asgi_app.py

from __future__ import annotations

import inspect
import json
import logging
import math
import secrets
import tempfile
import threading
from collections.abc import Callable
from pathlib import Path
from typing import Any, get_args, get_origin
from urllib.parse import unquote

from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import FileResponse, JSONResponse, Response, StreamingResponse
from starlette.routing import Route

from trackio import utils
from trackio.exceptions import TrackioAPIError
from trackio.remote_client import HTTP_API_VERSION

logger = logging.getLogger("trackio.asgi_app")

_PACKAGE_JSON_PATH = Path(__file__).parent / "package.json"
_TRACKIO_PACKAGE_VERSION = json.loads(_PACKAGE_JSON_PATH.read_text())["version"]


def _normalize_allowed_file_roots(
    allowed_file_roots: list[str | Path] | None,
) -> tuple[Path, ...]:
    roots = []
    for root in allowed_file_roots or []:
        roots.append(Path(root).resolve())
    return tuple(roots)


def _is_allowed_file_path(path: Path, allowed_roots: tuple[Path, ...]) -> bool:
    resolved_path = path.resolve(strict=False)
    for root in allowed_roots:
        try:
            resolved_path.relative_to(root)
            return True
        except ValueError:
            continue
    return False


def _json_safe(data: Any) -> Any:
    if data is None or isinstance(data, (str, bool, int)):
        return data
    if isinstance(data, float):
        return data if math.isfinite(data) else None
    if isinstance(data, dict):
        return {k: _json_safe(v) for k, v in data.items()}
    if isinstance(data, (list, tuple)):
        return [_json_safe(v) for v in data]
    if hasattr(data, "item"):
        try:
            return _json_safe(data.item())
        except Exception:
            pass
    return str(data)


def register_uploaded_temp_file(request: Request, file_path: str | Path) -> None:
    resolved_path = Path(file_path).resolve(strict=False)
    with request.app.state.uploaded_temp_files_lock:
        request.app.state.uploaded_temp_files.add(resolved_path)


def consume_uploaded_temp_file(request: Request, file_data: Any) -> Path:
    file_path = file_data.get("path") if isinstance(file_data, dict) else None
    if not isinstance(file_path, str) or not file_path:
        raise TrackioAPIError("Expected uploaded file metadata with a valid path.")

    resolved_path = Path(file_path).resolve(strict=False)
    with request.app.state.uploaded_temp_files_lock:
        if resolved_path not in request.app.state.uploaded_temp_files:
            raise TrackioAPIError(
                "Uploaded file was not created by this Trackio server."
            )
        request.app.state.uploaded_temp_files.remove(resolved_path)

    if not resolved_path.is_file():
        raise TrackioAPIError("Uploaded file is missing.")

    return resolved_path


def cleanup_uploaded_temp_file(file_path: str | Path) -> None:
    try:
        Path(file_path).unlink(missing_ok=True)
    except Exception:
        pass


def _invoke_handler(
    fn: Any,
    request: Request,
    args: list[Any] | None = None,
    kwargs: dict[str, Any] | None = None,
) -> Any:
    sig = inspect.signature(fn)
    params = list(sig.parameters.values())
    positional_args: list[Any] = []
    keyword_args: dict[str, Any] = {}
    args = list(args or [])
    kwargs = dict(kwargs or {})
    data_index = 0

    for param in params:
        if param.name == "request":
            keyword_args["request"] = request
        elif param.kind == inspect.Parameter.VAR_POSITIONAL:
            positional_args.extend(args[data_index:])
            data_index = len(args)
        elif param.kind == inspect.Parameter.VAR_KEYWORD:
            keyword_args.update(kwargs)
            kwargs.clear()
        elif param.name in kwargs:
            keyword_args[param.name] = kwargs.pop(param.name)
        elif param.kind in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        ) and data_index < len(args):
            positional_args.append(args[data_index])
            data_index += 1
        elif param.default is inspect.Signature.empty and param.kind not in (
            inspect.Parameter.VAR_POSITIONAL,
            inspect.Parameter.VAR_KEYWORD,
        ):
            raise TrackioAPIError(f"Missing required parameter: {param.name}")

    return fn(*positional_args, **keyword_args)


async def version_handler(request: Request) -> Response:
    mcp_enabled = bool(getattr(request.app.state, "mcp_enabled", False))
    return JSONResponse(
        {
            "version": _TRACKIO_PACKAGE_VERSION,
            "api_version": HTTP_API_VERSION,
            "api_transport": "http",
            "mcp_enabled": mcp_enabled,
            "mcp_path": "/mcp" if mcp_enabled else None,
        }
    )


def _json_schema_and_python_type(annotation: Any) -> tuple[dict[str, Any], str]:
    if annotation is inspect.Parameter.empty:
        return {"type": "object"}, "Any"
    origin = get_origin(annotation)
    args = get_args(annotation)
    if origin is not None and args:
        non_none = tuple(a for a in args if a is not type(None))
        if len(non_none) == 1 and len(args) > 1:
            return _json_schema_and_python_type(non_none[0])
        if origin is list:
            inner, py_inner = _json_schema_and_python_type(
                args[0] if args else inspect.Parameter.empty
            )
            return {"type": "array", "items": inner}, f"list[{py_inner}]"
        if origin is dict:
            return {"type": "object"}, "dict"
    if annotation in (str, bytes):
        return {"type": "string"}, "str"
    if annotation is int:
        return {"type": "integer"}, "int"
    if annotation is float:
        return {"type": "number"}, "float"
    if annotation is bool:
        return {"type": "boolean"}, "bool"
    name = getattr(annotation, "__name__", None)
    if name:
        return {"type": "object"}, name
    return {"type": "object"}, "Any"


def build_gradio_api_info(api_registry: dict[str, Any]) -> dict[str, Any]:
    named_endpoints: dict[str, Any] = {}
    for name in sorted(api_registry.keys()):
        fn = api_registry[name]
        if not callable(fn):
            continue
        sig = inspect.signature(fn)
        parameters: list[dict[str, Any]] = []
        for pname, param in sig.parameters.items():
            if pname == "request":
                continue
            jtype, pytype = _json_schema_and_python_type(param.annotation)
            has_default = param.default is not inspect.Parameter.empty
            parameters.append(
                {
                    "label": pname,
                    "parameter_name": pname,
                    "parameter_has_default": has_default,
                    "parameter_default": None if not has_default else param.default,
                    "type": jtype,
                    "python_type": {"type": pytype, "description": ""},
                    "component": "Api",
                    "example_input": None,
                }
            )
        ret_ann = sig.return_annotation
        if ret_ann is inspect.Signature.empty:
            ret_ann = Any
        rjtype, rpytype = _json_schema_and_python_type(ret_ann)
        returns = [
            {
                "label": "result",
                "type": rjtype,
                "python_type": {"type": rpytype, "description": ""},
                "component": "Api",
            }
        ]
        named_endpoints[f"/{name}"] = {
            "parameters": parameters,
            "returns": returns,
            "api_visibility": "public",
        }
    return {"named_endpoints": named_endpoints, "unnamed_endpoints": {}}


_MAX_GRADIO_CALL_EVENTS = 256


def _hf_token_value_is_unset(value: Any) -> bool:
    if value is None:
        return True
    if isinstance(value, str) and value.strip() == "":
        return True
    return False


def _authorization_bearer_token(request: Request) -> str | None:
    auth = request.headers.get("authorization") or request.headers.get("Authorization")
    if not auth:
        return None
    parts = auth.split()
    if len(parts) != 2 or parts[0].lower() != "bearer":
        return None
    tok = parts[1].strip()
    return tok or None


def _maybe_apply_hf_token_from_authorization(
    request: Request, fn: Any, args: list[Any], kwargs: dict[str, Any]
) -> None:
    if not utils.on_spaces():
        return
    token = _authorization_bearer_token(request)
    if not token:
        return
    sig = inspect.signature(fn)
    if "hf_token" not in sig.parameters:
        return
    params = [p for p in sig.parameters.values() if p.name != "request"]
    names = [p.name for p in params]
    if "hf_token" not in names:
        return
    idx = names.index("hf_token")
    if "hf_token" in kwargs:
        if _hf_token_value_is_unset(kwargs["hf_token"]):
            kwargs["hf_token"] = token
        return
    if idx < len(args):
        if _hf_token_value_is_unset(args[idx]):
            args[idx] = token
        return
    kwargs["hf_token"] = token


def _store_gradio_call_result(
    request: Request, event_id: str, api_name: str, data: Any
) -> None:
    with request.app.state.gradio_call_events_lock:
        d = request.app.state.gradio_call_events
        while len(d) >= _MAX_GRADIO_CALL_EVENTS:
            d.pop(next(iter(d)))
        d[event_id] = {"api_name": api_name, "data": data}


async def run_api_request(request: Request, api_name: str) -> Response:
    api_registry = request.app.state.api_registry
    fn = api_registry.get(api_name)
    if fn is None:
        return JSONResponse({"error": f"Unknown API: {api_name}"}, status_code=404)

    try:
        body = await request.json()
    except Exception:
        body = {}

    args: list[Any] = []
    kwargs: dict[str, Any] = {}
    if isinstance(body, dict):
        if "args" in body or "kwargs" in body:
            args = body.get("args") or []
            kwargs = body.get("kwargs") or {}
        elif "data" in body and isinstance(body["data"], list):
            args = body["data"]
        else:
            kwargs = body
    elif isinstance(body, list):
        args = body
    elif body is not None:
        args = [body]

    if not isinstance(args, list):
        args = [args]
    if not isinstance(kwargs, dict):
        kwargs = {}

    _maybe_apply_hf_token_from_authorization(request, fn, args, kwargs)

    try:
        result = _invoke_handler(fn, request, args=args, kwargs=kwargs)
        return JSONResponse({"data": _json_safe(result)})
    except TrackioAPIError as e:
        return JSONResponse({"error": str(e)}, status_code=400)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


async def api_handler(request: Request) -> Response:
    return await run_api_request(request, request.path_params["api_name"])


async def gradio_api_info_handler(request: Request) -> Response:
    return JSONResponse(
        build_gradio_api_info(request.app.state.api_registry),
        headers={"Cache-Control": "no-store"},
    )


async def gradio_call_post_handler(request: Request) -> Response:
    api_name = request.path_params["api_name"]
    resp = await run_api_request(request, api_name)
    if resp.status_code != 200:
        return resp
    body = json.loads(bytes(resp.body).decode())
    event_id = secrets.token_urlsafe(16)
    _store_gradio_call_result(request, event_id, api_name, body["data"])
    return JSONResponse({"event_id": event_id})


async def gradio_call_poll_handler(request: Request) -> Response:
    api_name = request.path_params["api_name"]
    event_id = request.path_params["event_id"]
    with request.app.state.gradio_call_events_lock:
        event = request.app.state.gradio_call_events.pop(event_id, None)
    if event is None:
        logger.info("gradio_api poll: unknown or expired event_id")
        return JSONResponse({"error": "Unknown or expired event_id"}, status_code=404)
    if event.get("api_name") != api_name:
        logger.info(
            "gradio_api poll: api_name mismatch (path=%r, stored=%r)",
            api_name,
            event.get("api_name"),
        )
        with request.app.state.gradio_call_events_lock:
            d = request.app.state.gradio_call_events
            while len(d) >= _MAX_GRADIO_CALL_EVENTS:
                d.pop(next(iter(d)))
            d[event_id] = event
        return JSONResponse({"error": "Unknown or expired event_id"}, status_code=404)

    data = event["data"]
    payload = json.dumps(_json_safe([data]))

    async def sse() -> Any:
        yield f"event: complete\ndata: {payload}\n\n"

    return StreamingResponse(
        sse(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-store",
            "X-Accel-Buffering": "no",
        },
    )


async def upload_handler(request: Request) -> Response:
    upload_authorizer = getattr(request.app.state, "upload_authorizer", None)
    if callable(upload_authorizer):
        try:
            upload_authorizer(request)
        except TrackioAPIError as e:
            return JSONResponse({"error": str(e)}, status_code=400)

    form = await request.form()
    uploads = form.getlist("files")
    saved_paths = []
    for upload in uploads:
        suffix = Path(getattr(upload, "filename", "") or "").suffix
        with tempfile.NamedTemporaryFile(
            delete=False,
            prefix="trackio-upload-",
            suffix=suffix,
        ) as tmp:
            tmp.write(await upload.read())
            register_uploaded_temp_file(request, tmp.name)
            saved_paths.append(tmp.name)
    return JSONResponse({"paths": saved_paths})


async def gradio_upload_alias_handler(request: Request) -> Response:
    return await upload_handler(request)


_DISALLOWED_FILE_SUFFIXES = frozenset(
    {".db", ".db-journal", ".db-wal", ".db-shm", ".sqlite", ".sqlite3"}
)


async def file_handler(request: Request) -> Response:
    fs_path = request.query_params.get("path")
    if fs_path is None:
        return Response("Missing path", status_code=400)
    fp = Path(unquote(fs_path)).resolve(strict=False)
    if fp.suffix.lower() in _DISALLOWED_FILE_SUFFIXES:
        return Response("Not found", status_code=404)
    allowed_roots = getattr(request.app.state, "allowed_file_roots", ())
    if fp.is_file() and _is_allowed_file_path(fp, allowed_roots):
        return FileResponse(str(fp))
    return Response("Not found", status_code=404)


async def artifact_blob_handler(request: Request) -> Response:
    from trackio import cas  # noqa: PLC0415
    from trackio import server as _server  # noqa: PLC0415

    project = request.path_params.get("project")
    digest = request.path_params.get("digest")
    try:
        project = _server._validate_project_name(project)
        digest = _server._validate_sha256_digest(digest)
    except TrackioAPIError:
        return Response("Not found", status_code=404)
    blob = cas.blob_path(project, digest)
    if not blob.is_file():
        return Response("Not found", status_code=404)
    return FileResponse(str(blob))


def create_trackio_starlette_app(
    oauth_routes: list[Route],
    api_registry: dict[str, Any],
    extra_routes: list[Any] | None = None,
    mcp_lifespan: Any = None,
    mcp_enabled: bool = False,
    allowed_file_roots: list[str | Path] | None = None,
    upload_authorizer: Callable[[Request], None] | None = None,
) -> Starlette:
    routes: list[Any] = list(oauth_routes)
    routes.extend(
        [
            Route("/version", endpoint=version_handler, methods=["GET"]),
            Route("/api/upload", endpoint=upload_handler, methods=["POST"]),
            Route("/api/{api_name:str}", endpoint=api_handler, methods=["POST"]),
            Route("/file", endpoint=file_handler, methods=["GET"]),
            Route(
                "/artifact_blob/{project:str}/{digest:str}",
                endpoint=artifact_blob_handler,
                methods=["GET"],
            ),
        ]
    )
    if utils.on_spaces():
        routes.extend(
            [
                Route(
                    "/gradio_api/info",
                    endpoint=gradio_api_info_handler,
                    methods=["GET"],
                ),
                Route(
                    "/gradio_api/info/",
                    endpoint=gradio_api_info_handler,
                    methods=["GET"],
                ),
                Route(
                    "/gradio_api/upload",
                    endpoint=gradio_upload_alias_handler,
                    methods=["POST"],
                ),
                Route(
                    "/gradio_api/upload/",
                    endpoint=gradio_upload_alias_handler,
                    methods=["POST"],
                ),
                Route(
                    "/gradio_api/call/{api_name:str}",
                    endpoint=gradio_call_post_handler,
                    methods=["POST"],
                ),
                Route(
                    "/gradio_api/call/{api_name:str}/",
                    endpoint=gradio_call_post_handler,
                    methods=["POST"],
                ),
                Route(
                    "/gradio_api/call/{api_name:str}/{event_id:str}",
                    endpoint=gradio_call_poll_handler,
                    methods=["GET"],
                ),
                Route(
                    "/gradio_api/call/{api_name:str}/{event_id:str}/",
                    endpoint=gradio_call_poll_handler,
                    methods=["GET"],
                ),
            ]
        )
    routes.extend(extra_routes or [])
    app = Starlette(routes=routes, lifespan=mcp_lifespan)
    app.state.api_registry = api_registry
    app.state.mcp_enabled = mcp_enabled
    app.state.allowed_file_roots = _normalize_allowed_file_roots(allowed_file_roots)
    app.state.upload_authorizer = upload_authorizer
    app.state.uploaded_temp_files = set()
    app.state.uploaded_temp_files_lock = threading.Lock()
    if utils.on_spaces():
        app.state.gradio_call_events = {}
        app.state.gradio_call_events_lock = threading.Lock()
    return app