invoke-ai--invokeai/invokeai/app/api/routers/download_queue.py

# Copyright (c) 2023 Lincoln D. Stein
"""FastAPI route for the download queue."""

from pathlib import Path as FsPath
from pathlib import PurePosixPath, PureWindowsPath
from typing import List, Optional

from fastapi import Body, Path, Response
from fastapi.routing import APIRouter
from pydantic.networks import AnyHttpUrl
from starlette.exceptions import HTTPException

from invokeai.app.api.auth_dependencies import AdminUserOrDefault, CurrentUserOrDefault
from invokeai.app.api.dependencies import ApiDependencies
from invokeai.app.services.download import (
    DownloadJob,
    UnknownJobIDException,
)

download_queue_router = APIRouter(prefix="/v1/download_queue", tags=["download_queue"])


def _validate_dest(dest: str) -> str:
    """Reject absolute paths and parent-traversal segments.

    Accepts a relative POSIX- or Windows-style path. Returns the original string
    for the caller to wrap in `Path(...)`. Raises 400 on suspicious input so the
    download service never sees it.
    """
    if not dest or not dest.strip():
        raise HTTPException(status_code=400, detail="Download destination must not be empty.")

    posix = PurePosixPath(dest)
    windows = PureWindowsPath(dest)
    if posix.is_absolute() or windows.is_absolute():
        raise HTTPException(status_code=400, detail="Download destination must be a relative path.")

    if ".." in posix.parts or ".." in windows.parts:
        raise HTTPException(status_code=400, detail="Download destination must not contain '..' segments.")

    return dest


@download_queue_router.get(
    "/",
    operation_id="list_downloads",
)
async def list_downloads(current_user: CurrentUserOrDefault) -> List[DownloadJob]:
    """Get a list of active and inactive jobs."""
    queue = ApiDependencies.invoker.services.download_queue
    return queue.list_jobs()


@download_queue_router.patch(
    "/",
    operation_id="prune_downloads",
    responses={
        204: {"description": "All completed jobs have been pruned"},
        400: {"description": "Bad request"},
    },
)
async def prune_downloads(current_user: AdminUserOrDefault) -> Response:
    """Prune completed and errored jobs."""
    queue = ApiDependencies.invoker.services.download_queue
    queue.prune_jobs()
    return Response(status_code=204)


@download_queue_router.post(
    "/i/",
    operation_id="download",
)
async def download(
    current_user: CurrentUserOrDefault,
    source: AnyHttpUrl = Body(description="download source"),
    dest: str = Body(description="download destination"),
    priority: int = Body(default=10, description="queue priority"),
    access_token: Optional[str] = Body(default=None, description="token for authorization to download"),
) -> DownloadJob:
    """Download the source URL to the file or directory indicted in dest."""
    validated_dest = _validate_dest(dest)
    queue = ApiDependencies.invoker.services.download_queue
    return queue.download(source, FsPath(validated_dest), priority, access_token)


@download_queue_router.get(
    "/i/{id}",
    operation_id="get_download_job",
    responses={
        200: {"description": "Success"},
        404: {"description": "The requested download JobID could not be found"},
    },
)
async def get_download_job(
    current_user: CurrentUserOrDefault,
    id: int = Path(description="ID of the download job to fetch."),
) -> DownloadJob:
    """Get a download job using its ID."""
    try:
        job = ApiDependencies.invoker.services.download_queue.id_to_job(id)
        return job
    except UnknownJobIDException as e:
        raise HTTPException(status_code=404, detail=str(e))


@download_queue_router.delete(
    "/i/{id}",
    operation_id="cancel_download_job",
    responses={
        204: {"description": "Job has been cancelled"},
        404: {"description": "The requested download JobID could not be found"},
    },
)
async def cancel_download_job(
    current_user: CurrentUserOrDefault,
    id: int = Path(description="ID of the download job to cancel."),
) -> Response:
    """Cancel a download job using its ID."""
    try:
        queue = ApiDependencies.invoker.services.download_queue
        job = queue.id_to_job(id)
        queue.cancel_job(job)
        return Response(status_code=204)
    except UnknownJobIDException as e:
        raise HTTPException(status_code=404, detail=str(e))


@download_queue_router.delete(
    "/i",
    operation_id="cancel_all_download_jobs",
    responses={
        204: {"description": "Download jobs have been cancelled"},
    },
)
async def cancel_all_download_jobs(current_user: AdminUserOrDefault) -> Response:
    """Cancel all download jobs."""
    ApiDependencies.invoker.services.download_queue.cancel_all_jobs()
    return Response(status_code=204)