fastapi_worker.py

# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Custom NAT FastAPI worker with this workflow's server-side customizations.

This is the single umbrella worker wired in via the FastAPI front-end config
(``runner_class: vuln_analysis.server.fastapi_worker.VulnAnalysisFastApiWorker``). It
subclasses NAT's default FastAPI worker and layers on features specific to this workflow —
each safe to leave configured because it is either read-only or inert until enabled by env:

- **Read-only report endpoints** (always on). The pipeline agents write Markdown reports
  (and a JSON sidecar for VEX results) to disk under ``output_dir``; after wiring up all of
  NAT's built-in routes this worker adds:

  - ``GET /reports`` — list available reports organized by group and report type. Accepts an
    optional ``?container=``, ``?vulnerability=``, and ``?report_type=`` filters. Known
    CVE/GHSA aliases resolve to the same report entries.
  - ``GET /reports/{group}/{report_type}/{filename}`` — fetch one report inline (for the UI
    to render) or as a download (``?download=true``).

  Reports are only served from a single configured reports-root (path-traversal safe). The
  root defaults to the workflow's ``output_dir`` and is overridable via the
  ``VULN_ANALYSIS_REPORTS_ROOT`` env var (useful for browsing reports produced by a prior
  ``nat run``).

- **Provider-aware preflight** (always on). ``POST /preflight`` validates a pipeline input,
  local prerequisites, and credentials through lightweight provider readiness endpoints.
  It returns HTTP 503 with per-check remediations when the run is not ready.

- **Optional inbound OAuth2/JWT bearer-token auth** (inert unless enabled). When
  ``VULN_ANALYSIS_AUTH_ENABLED=true`` (see :class:`~vuln_analysis.server.auth.AuthConfig`),
  ``build_app`` installs ``JwtAuthMiddleware`` after NAT builds the app. When disabled the
  custom routes remain available without authentication.
"""

import asyncio
import json
import logging
import os
from datetime import UTC
from datetime import datetime
from pathlib import Path
from typing import Literal

from fastapi import FastAPI
from fastapi import HTTPException
from fastapi import Query
from fastapi import Request
from fastapi.responses import FileResponse
from fastapi.responses import JSONResponse
from fastapi.responses import Response
from fastapi.routing import APIRoute
from nat.front_ends.fastapi.async_jobs.job_store import JobStatus
from nat.front_ends.fastapi.fastapi_front_end_config import AsyncGenerationStatusResponse
from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorker
from pydantic import BaseModel

from vuln_analysis.data_models.input import PipelineInput
from vuln_analysis.data_models.output import PipelineResult
from vuln_analysis.pipeline.helpers import REPORT_SUBDIRS
from vuln_analysis.pipeline.helpers import split_report_stem
from vuln_analysis.pipeline.helpers import to_report_url
from vuln_analysis.preflight import PreflightResult
from vuln_analysis.preflight import run_live_preflight
from vuln_analysis.semantic_failure import SemanticWorkflowError
from vuln_analysis.semantic_failure import is_semantic_failure
from vuln_analysis.semantic_failure import require_successful_pipeline_result
from vuln_analysis.semantic_failure import semantic_failure_message
from vuln_analysis.server.async_job_runner import install_async_job_runner
from vuln_analysis.server.auth import AuthConfig
from vuln_analysis.server.auth import JwtAuthMiddleware
from vuln_analysis.server.job_progress import ExtractionProgress
from vuln_analysis.server.job_progress import JobProgressStore
from vuln_analysis.server.webhooks import WebhookConfig
from vuln_analysis.server.webhooks import start_webhook_dispatcher
from vuln_analysis.server_identity import DOCUMENTATION_VALIDATION_RUN_ID_ENV

logger = logging.getLogger(__name__)

# Fields in a VEX JSON sidecar that hold absolute server filesystem paths; rewritten to
# endpoint-relative URLs before the sidecar is served so we don't leak server paths.
_SIDECAR_PATH_FIELDS = ("report_path", "exploitability_report_path")

# Summary fields surfaced from a VEX sidecar in the /reports listing.
_VEX_SUMMARY_FIELDS = (
    "cve_id",
    "container",
    "status",
    "justification",
    "exploitability_verdict",
    "confidence",
    "risk_level",
)

_VULNERABILITY_ALIAS_DIR = ".cache_aliases"
_REPORTS_COLLECTION_HEADERS = {"Link": '</reports>; rel="collection"'}

_READINESS_TIMEOUT_SECONDS = 3.0
_INTERRUPTED_JOB_ERROR = "Analysis service restarted before the job completed"


class VulnAnalysisAsyncGenerationStatusResponse(AsyncGenerationStatusResponse):
    """NAT async status plus durable extraction and workflow-stage progress."""

    progress: ExtractionProgress | None = None


class WorkflowErrorResponse(BaseModel):
    """Typed API envelope for a workflow result that failed its terminal contract."""

    code: Literal["workflow_error"] = "workflow_error"
    message: str
    details: str | None = None


def _utc(value: datetime) -> datetime:
    return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC)


def _latest_timestamp(job_updated_at: datetime, progress: ExtractionProgress | None) -> datetime:
    if progress is None:
        return job_updated_at
    return max(_utc(job_updated_at), _utc(progress.updated_at))


def replace_sync_generation_routes(app: FastAPI) -> int:
    """Return HTTP 422 when a synchronous pipeline result reports ``success:false``."""
    replaced: list[APIRoute] = []
    for route in list(app.routes):
        if (
            isinstance(route, APIRoute)
            and route.response_model is PipelineResult
            and route.methods == {"POST"}
        ):
            app.routes.remove(route)
            replaced.append(route)

    def build_semantic_endpoint(original: APIRoute):
        async def post_generate(payload: PipelineInput, request: Request, response: Response):
            result = await original.endpoint(response=response, request=request, payload=payload)
            if isinstance(result, Response):
                return result

            pipeline_result = PipelineResult.model_validate(result)
            if pipeline_result.outcome == "complete":
                try:
                    require_successful_pipeline_result(pipeline_result)
                except SemanticWorkflowError as exc:
                    logger.warning("Synchronous pipeline returned an unusable complete result: %s", exc)
                    workflow_error = WorkflowErrorResponse(
                        message="Pipeline returned an unusable complete result",
                    )
                    return JSONResponse(
                        status_code=422,
                        content=workflow_error.model_dump(mode="json", exclude_none=True),
                        headers={
                            "X-Workflow-Success": "false",
                            "X-Workflow-Outcome": "failed",
                        },
                    )
            response.headers["X-Workflow-Success"] = str(pipeline_result.success).lower()
            response.headers["X-Workflow-Outcome"] = pipeline_result.outcome
            if is_semantic_failure(pipeline_result):
                response.status_code = 422
                logger.warning(semantic_failure_message(pipeline_result))
            return pipeline_result

        return post_generate

    for original in replaced:
        responses = dict(original.responses or {})
        responses[422] = {
            "model": PipelineResult | WorkflowErrorResponse,
            "description": (
                "The request was valid, but no requested vulnerability produced a usable result. "
                "The response body remains a PipelineResult with outcome=failed. Partial results "
                "return HTTP 200 with outcome=partial. A result that declares outcome=complete "
                "but fails the terminal artifact contract returns a workflow_error envelope."
            )
        }
        app.add_api_route(
            path=original.path,
            endpoint=build_semantic_endpoint(original),
            methods=["POST"],
            response_model=PipelineResult,
            name=original.name,
            tags=original.tags,
            summary=original.summary,
            description=original.description,
            responses=responses,
            deprecated=original.deprecated,
            include_in_schema=original.include_in_schema,
        )

    if replaced:
        logger.info("Added semantic failure propagation to %d synchronous route(s)", len(replaced))
    return len(replaced)


def replace_async_generation_status_routes(
    app: FastAPI,
    _job_store,
    progress_store: JobProgressStore,
) -> int:
    """Replace NAT generation-status GET routes with response-compatible enriched routes."""
    replaced: list[APIRoute] = []
    for route in list(app.routes):
        if (
            isinstance(route, APIRoute)
            and route.response_model is AsyncGenerationStatusResponse
            and route.methods == {"GET"}
        ):
            app.routes.remove(route)
            replaced.append(route)

    def build_status_endpoint(original: APIRoute):
        async def get_status(job_id: str, http_request: Request):
            path = original.path
            logger.info("Getting enriched async status for job %s via %s", job_id, path)
            try:
                base_response = await original.endpoint(job_id, http_request)
            except HTTPException as exc:
                if exc.status_code == 404:
                    await progress_store.delete(job_id)
                raise

            if not isinstance(base_response, AsyncGenerationStatusResponse):
                base_response = AsyncGenerationStatusResponse.model_validate(base_response)
            progress = await progress_store.get(job_id)
            if (
                base_response.expires_at is not None
                and _utc(base_response.expires_at) <= datetime.now(UTC)
            ):
                await progress_store.delete(job_id)
                progress = None
            payload = base_response.model_dump()
            payload["progress"] = progress
            payload["updated_at"] = _latest_timestamp(base_response.updated_at, progress)
            return VulnAnalysisAsyncGenerationStatusResponse.model_validate(payload)

        return get_status

    for original in replaced:
        app.add_api_route(
            path=original.path,
            endpoint=build_status_endpoint(original),
            methods=["GET"],
            response_model=VulnAnalysisAsyncGenerationStatusResponse,
            name=original.name,
            tags=original.tags,
            summary=original.summary,
            description=(
                f"{original.description or 'Get the status of an async job.'} "
                "Active jobs include structured extraction and workflow-stage progress."
            ),
            responses=original.responses,
            deprecated=original.deprecated,
        )

    if replaced:
        logger.info(
            "Enriched %d async generation status route(s) with workflow progress",
            len(replaced),
        )
    return len(replaced)


def _vex_summary(sidecar: Path) -> dict:
    """Return selected summary fields from a VEX JSON sidecar (empty on any error)."""
    try:
        data = json.loads(sidecar.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}
    return {k: data[k] for k in _VEX_SUMMARY_FIELDS if k in data}


def _normalize_vulnerability(value: str) -> str:
    """Normalize a CVE/GHSA identifier for report discovery comparisons."""
    value = value.strip()
    if value.upper().startswith("CVE-"):
        return value.upper()
    if value.upper().startswith("GHSA-"):
        return f"GHSA-{value[5:].lower()}"
    return value


def _vulnerability_aliases(root: Path, group_dir: Path) -> dict[str, set[str]]:
    """Map every known CVE/GHSA alias to the complete identifier set for its report."""
    report_dir = group_dir / "vulnerability_reports"
    aliases_by_id: dict[str, set[str]] = {}
    if not report_dir.is_dir():
        return aliases_by_id

    for report_file in report_dir.glob("*.md"):
        if not _is_safe_report_file(root, report_file):
            continue
        identifiers = {_normalize_vulnerability(report_file.stem)}
        metadata_path = report_dir / _VULNERABILITY_ALIAS_DIR / f"{report_file.stem}.json"
        try:
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
            if not isinstance(metadata, dict):
                raise ValueError("alias metadata must be an object")
            metadata_report = metadata.get("report")
            metadata_identifiers = metadata.get("identifiers")
            if metadata_report != report_file.name:
                raise ValueError("alias metadata report does not match")
            if not isinstance(metadata_identifiers, list) or not all(
                    isinstance(item, str) for item in metadata_identifiers):
                raise ValueError("alias metadata identifiers must be strings")
            identifiers.update(_normalize_vulnerability(item) for item in metadata_identifiers)
        except (OSError, TypeError, ValueError):
            pass

        for identifier in identifiers:
            aliases_by_id.setdefault(identifier, set()).update(identifiers)

    return aliases_by_id


def _identifiers_for_report(
    report_type: str,
    stem: str,
    aliases_by_id: dict[str, set[str]],
) -> list[str]:
    """Return every known identifier that resolves to a report."""
    if report_type == "container_reports":
        return []
    identifier = stem if report_type == "vulnerability_reports" else split_report_stem(stem)[0]
    if not identifier:
        return []
    normalized = _normalize_vulnerability(identifier)
    return sorted(aliases_by_id.get(normalized, {normalized}))


def _report_container(report_type: str, stem: str) -> str | None:
    if report_type == "container_reports":
        return stem
    if report_type in {"exploitability_reports", "vex_results"}:
        return split_report_stem(stem)[1]
    return None


def _report_group(root: Path, report_file: Path) -> str:
    relative = report_file.parent.parent.resolve().relative_to(root)
    return relative.as_posix() if relative.parts else "default"


def _is_safe_report_file(root: Path, report_file: Path) -> bool:
    try:
        return report_file.is_file() and report_file.resolve().is_relative_to(root)
    except OSError:
        return False


def _build_entry(
    root: Path,
    report_type: str,
    report_file: Path,
    aliases_by_id: dict[str, set[str]],
) -> dict:
    """Build a /reports listing entry for a single report file."""
    try:
        size = report_file.stat().st_size
    except OSError:
        # File may have been removed or become inaccessible between the
        # is_file() filter and here (TOCTOU); fall back to an unknown size.
        size = None
    entry = {
        "filename": report_file.name,
        "size": size,
        "url": to_report_url(str(report_file), root),
        "group": _report_group(root, report_file),
        "report_type": report_type,
    }
    identifiers = _identifiers_for_report(report_type, report_file.stem, aliases_by_id)
    if identifiers:
        entry["vulnerability_ids"] = identifiers
    container = _report_container(report_type, report_file.stem)
    if container:
        entry["container"] = container
    if report_type == "vex_results":
        sidecar = report_file.with_suffix(".json")
        if sidecar.is_file():
            entry["json_url"] = to_report_url(str(sidecar), root)
            entry.update(_vex_summary(sidecar))
    return entry


def _normalize_container(value: str) -> str:
    """Normalize a container filter value to the pipeline's derived container name.

    Accepts either the derived name (``morpheus_25.06-runtime``) or a full/partial image
    reference (``nvcr.io/nvidia/morpheus/morpheus:25.06-runtime``, ``morpheus:25.06-runtime``).
    A reference with a tag is converted to ``<image_basename>_<tag>`` to match on-disk report
    filenames (see ``derive_report_filename``). Registry prefixes, ports, and digests are removed
    before matching an untagged image basename.
    """
    value = value.strip()
    last_segment = value.rsplit("/", 1)[-1]
    image_segment = last_segment.split("@", 1)[0]
    if ":" in image_segment:
        base, _, tag = image_segment.rpartition(":")
        return f"{base}_{tag}"
    return image_segment


def _matched_vulnerability_ids(
    root: Path,
    group_dir: Path,
    container: str,
    aliases_by_id: dict[str, set[str]],
) -> set[str]:
    """Vulnerability IDs and aliases that have a per-container report."""
    identifiers: set[str] = set()
    for report_type in ("exploitability_reports", "vex_results"):
        type_dir = group_dir / report_type
        if not type_dir.is_dir():
            continue
        for report_file in type_dir.glob("*.md"):
            if not _is_safe_report_file(root, report_file):
                continue
            vuln_id, cont = split_report_stem(report_file.stem)
            if vuln_id and cont == container:
                normalized = _normalize_vulnerability(vuln_id)
                identifiers.update(aliases_by_id.get(normalized, {normalized}))
    return identifiers


def _matches_container(
    report_type: str,
    stem: str,
    container: str,
    matched_vulnerabilities: set[str],
) -> bool:
    """Whether a report belongs to *container*.

    Container reports match by name; exploitability/VEX reports match by their container
    suffix; vulnerability reports (container-independent) match via the CVEs that have a
    report for the container.
    """
    if report_type == "container_reports":
        return stem == container
    if report_type == "vulnerability_reports":
        return _normalize_vulnerability(stem) in matched_vulnerabilities
    return split_report_stem(stem)[1] == container


def _matches_vulnerability(
    report_type: str,
    stem: str,
    vulnerability: str,
    aliases_by_id: dict[str, set[str]],
) -> bool:
    """Whether a vulnerability-scoped report carries a CVE/GHSA identifier or alias."""
    if report_type == "container_reports":
        return False
    identifiers = _identifiers_for_report(report_type, stem, aliases_by_id)
    return vulnerability in identifiers


def _collect_group_reports(
    root: Path,
    group_dir: Path,
    container: str | None = None,
    vulnerability: str | None = None,
    report_type_filter: str | None = None,
) -> dict:
    """Collect Markdown reports under *group_dir*, grouped by report type.

    When *container* is given, only reports for that container are included.
    When *vulnerability* is given, image-wide container reports are excluded.
    """
    aliases_by_id = _vulnerability_aliases(root, group_dir)
    matched_vulnerabilities = (
        _matched_vulnerability_ids(root, group_dir, container, aliases_by_id)
        if container is not None
        else set()
    )
    reports: dict[str, list[dict]] = {}
    for report_type in REPORT_SUBDIRS:
        if report_type_filter is not None and report_type != report_type_filter:
            continue
        type_dir = group_dir / report_type
        if not type_dir.is_dir():
            continue
        entries = [
            _build_entry(root, report_type, report_file, aliases_by_id)
            for report_file in sorted(type_dir.glob("*.md"))
            if _is_safe_report_file(root, report_file)
            and (
                container is None
                or _matches_container(
                    report_type,
                    report_file.stem,
                    container,
                    matched_vulnerabilities,
                )
            )
            and (
                vulnerability is None
                or _matches_vulnerability(
                    report_type,
                    report_file.stem,
                    vulnerability,
                    aliases_by_id,
                )
            )
        ]
        if entries:
            reports[report_type] = entries
    return reports


def _list_reports(
    root: Path,
    container: str | None = None,
    vulnerability: str | None = None,
    report_type: str | None = None,
) -> dict:
    """Enumerate reports as ``{"groups": [{"group", "reports": {<type>: [...]}}]}``."""
    groups: list[dict] = []
    if not root.is_dir():
        return {"groups": groups}

    # Implicit "default" group: report subdirs directly under the root (flat layout).
    default_reports = _collect_group_reports(root, root, container, vulnerability, report_type)
    if default_reports:
        groups.append({"group": "default", "reports": default_reports})

    # Named groups: immediate subdirectories that group their own report subdirs.
    for child in sorted(p for p in root.iterdir() if p.is_dir()):
        if child.name in REPORT_SUBDIRS:
            continue  # part of the default/flat layout, already collected above
        child_reports = _collect_group_reports(root, child, container, vulnerability, report_type)
        if child_reports:
            groups.append({"group": child.name, "reports": child_reports})

    return {"groups": groups}


def _sanitize_sidecar(path: Path, root: Path) -> str:
    """Return sidecar JSON with absolute path fields rewritten to relative report URLs."""
    try:
        raw = path.read_text(encoding="utf-8")
    except OSError as exc:
        # File may have been removed or become inaccessible between the caller's
        # is_file() check and this read (TOCTOU); surface a clean 404.
        raise HTTPException(
            status_code=404,
            detail="Report not found",
            headers=_REPORTS_COLLECTION_HEADERS,
        ) from exc
    try:
        data = json.loads(raw)
    except ValueError:
        return raw
    for field in _SIDECAR_PATH_FIELDS:
        if field in data:
            data[field] = to_report_url(data[field], root)
    return json.dumps(data, indent=2)


def _resolve_report_file(root: Path, group: str, report_type: str, filename: str) -> Path:
    """Resolve and validate a report file under *root*, or raise HTTPException."""
    if report_type not in REPORT_SUBDIRS:
        raise HTTPException(status_code=400, detail=f"Unknown report_type '{report_type}'")
    if not filename or filename in (".", "..") or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")

    base = root if group == "default" else root / group
    try:
        resolved = (base / report_type / filename).resolve()
    except OSError as exc:
        raise HTTPException(status_code=400, detail="Invalid report path") from exc
    if not resolved.is_relative_to(root):
        raise HTTPException(status_code=400, detail="Invalid report path")
    if not resolved.is_file():
        raise HTTPException(
            status_code=404,
            detail="Report not found",
            headers=_REPORTS_COLLECTION_HEADERS,
        )
    return resolved


def add_reports_routes(app: FastAPI, root: Path) -> None:
    """Register the report listing and fetch endpoints on *app*, scoped to *root*."""

    def list_reports(
        container: str | None = Query(
            default=None,
            description="Filter to reports for a specific container. Accepts the derived name "
            "('morpheus_25.06-runtime') or a full image reference "
            "('nvcr.io/nvidia/morpheus/morpheus:25.06-runtime').",
        ),
        vulnerability: str | None = Query(
            default=None,
            description=(
                "Filter to vulnerability-specific reports carrying a CVE or GHSA identifier. "
                "Known CVE/GHSA aliases resolve to the same reports; image-wide container "
                "reports are excluded."
            ),
        ),
        report_type: str | None = Query(
            default=None,
            description=f"Filter to one report type: {', '.join(REPORT_SUBDIRS)}.",
        ),
    ) -> dict:
        """List available reports organized by group and report type."""
        if report_type is not None and report_type not in REPORT_SUBDIRS:
            raise HTTPException(status_code=400, detail=f"Unknown report_type '{report_type}'")
        normalized = _normalize_container(container) if container else None
        normalized_vulnerability = _normalize_vulnerability(vulnerability) if vulnerability else None
        return _list_reports(
            root,
            container=normalized,
            vulnerability=normalized_vulnerability,
            report_type=report_type,
        )

    def get_report(
        group: str,
        report_type: str,
        filename: str,
        download: bool = Query(default=False, description="Force a file download"),
    ):
        """Fetch a single report (Markdown inline by default, or as an attachment)."""
        resolved = _resolve_report_file(root, group, report_type, filename)
        disposition = "attachment" if download else "inline"
        # Escape any double-quote in the filename per RFC 6266 so the header stays well-formed.
        escaped_name = resolved.name.replace('"', '\\"')
        headers = {"Content-Disposition": f'{disposition}; filename="{escaped_name}"'}

        if resolved.suffix == ".json":
            # Rewrite absolute path fields so served JSON doesn't leak server paths.
            return Response(
                content=_sanitize_sidecar(resolved, root),
                media_type="application/json",
                headers=headers,
            )
        media_type = "text/markdown; charset=utf-8" if resolved.suffix == ".md" else "application/octet-stream"
        return FileResponse(resolved, media_type=media_type, headers=headers)

    app.add_api_route(
        path="/reports",
        endpoint=list_reports,
        methods=["GET"],
        tags=["Reports"],
        description="List available agent reports organized by group and report type. "
        "Filter with ?container=<name>, ?vulnerability=<CVE-or-GHSA>, and/or "
        "?report_type=<type>. Known CVE/GHSA aliases resolve to the same reports; "
        "vulnerability filters exclude image-wide container reports.",
    )
    app.add_api_route(
        path="/reports/{group}/{report_type}/{filename}",
        endpoint=get_report,
        methods=["GET"],
        tags=["Reports"],
        description="Fetch a report (Markdown inline for viewing, or ?download=true to download).",
    )
    logger.info("Added report endpoints at /reports (root=%s)", root)


def add_readiness_route(app: FastAPI, job_store) -> None:
    """Register a readiness probe that requires a live Dask scheduler and worker."""

    async def readiness() -> dict:
        if job_store is None:
            raise HTTPException(status_code=503, detail="Dask job service is unavailable")

        try:
            scheduler_info = await asyncio.wait_for(
                asyncio.to_thread(lambda: job_store.dask_client.scheduler_info(n_workers=-1)),
                timeout=_READINESS_TIMEOUT_SECONDS,
            )
        except asyncio.TimeoutError as exc:
            raise HTTPException(status_code=503, detail="Dask scheduler readiness check timed out") from exc
        except Exception as exc:
            logger.warning("Dask scheduler readiness check failed: %s", exc)
            raise HTTPException(status_code=503, detail="Dask scheduler is unavailable") from exc

        worker_count = sum(worker.get("status") == "running" for worker in scheduler_info.get("workers", {}).values())
        if worker_count < 1:
            raise HTTPException(status_code=503, detail="Dask scheduler has no workers")

        response = {"status": "ready", "scheduler": "connected", "workers": worker_count}
        validation_run_id = os.environ.get(DOCUMENTATION_VALIDATION_RUN_ID_ENV)
        if validation_run_id:
            response["validation_run_id"] = validation_run_id
        return response

    app.add_api_route(
        path="/ready",
        endpoint=readiness,
        methods=["GET"],
        tags=["Health"],
        description="Readiness check for the analysis API, Dask scheduler, and workers",
    )
    logger.info("Added Dask-aware readiness endpoint at /ready")


async def reconcile_interrupted_jobs(
    job_store,
    *,
    scheduler_is_external: bool = False,
    progress_store: JobProgressStore | None = None,
) -> int:
    """Mark jobs left active by a prior local service process as interrupted."""
    if job_store is None:
        return 0
    if scheduler_is_external:
        logger.info("Preserving active jobs managed by an external Dask scheduler")
        return 0

    active_jobs = []
    for status in (JobStatus.SUBMITTED, JobStatus.RUNNING):
        active_jobs.extend(await job_store.get_jobs_by_status(status))

    interrupted_count = 0
    for job in active_jobs:
        try:
            await job_store.update_status(
                job.job_id,
                JobStatus.INTERRUPTED,
                error=_INTERRUPTED_JOB_ERROR,
            )
        except ValueError:
            logger.warning("Skipped restart reconciliation for missing job %s", job.job_id)
            continue
        if progress_store is not None:
            await progress_store.mark_terminal(
                job.job_id,
                "interrupted",
                _INTERRUPTED_JOB_ERROR,
            )
        interrupted_count += 1

    if interrupted_count:
        logger.warning("Marked %d job(s) interrupted after analysis-service restart", interrupted_count)
    return interrupted_count


def add_preflight_route(app: FastAPI) -> None:
    """Register the pipeline readiness endpoint."""

    async def preflight(
        inp: PipelineInput,
        response: Response,
    ) -> PreflightResult:
        result = await run_live_preflight(inp)
        if not result.ready:
            response.status_code = 503
        return result

    app.add_api_route(
        path="/preflight",
        endpoint=preflight,
        methods=["POST"],
        response_model=PreflightResult,
        tags=["Preflight"],
        description="Check credentials, input paths, and local tools without inference or search work.",
    )
    logger.info("Added provider-aware preflight endpoint at /preflight")


class VulnAnalysisFastApiWorker(FastApiFrontEndPluginWorker):
    """Default NAT FastAPI worker plus read-only report endpoints, optional JWT auth, and
    optional outbound job-completion webhooks."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._webhook_dispatcher = None
        self._progress_store = None

    def _reports_root(self) -> Path:
        override = os.environ.get("VULN_ANALYSIS_REPORTS_ROOT")
        if override:
            return Path(override).resolve()
        output_dir = getattr(self._config.workflow, "output_dir", "outputs")
        return Path(output_dir).resolve()

    async def add_routes(self, app: FastAPI, builder) -> None:
        install_async_job_runner()
        await super().add_routes(app, builder)
        replace_sync_generation_routes(app)
        if self._job_store is not None and self._db_url is not None:
            self._progress_store = JobProgressStore(self._db_url)
            await self._progress_store.initialize()
            removed = await self._progress_store.delete_expired_or_orphaned()
            if removed:
                logger.info("Removed %d expired or orphaned async progress record(s)", removed)
            replace_async_generation_status_routes(
                app,
                self._job_store,
                self._progress_store,
            )
        # NAT leaves this config value unset for its controller-owned LocalCluster. An
        # operator-supplied scheduler can outlive an API restart, so its jobs stay active.
        await reconcile_interrupted_jobs(
            self._job_store,
            scheduler_is_external=self.front_end_config.scheduler_address is not None,
            progress_store=self._progress_store,
        )
        add_readiness_route(app, self._job_store)
        add_preflight_route(app)
        add_reports_routes(app, self._reports_root())
        # Start the terminal-job webhook dispatcher. Inert unless the feature is enabled and
        # configured; returns None otherwise (see start_webhook_dispatcher).
        self._webhook_dispatcher = start_webhook_dispatcher(
            WebhookConfig.from_env(), self._job_store, self._reports_root()
        )

    async def cleanup_session_managers(self) -> None:
        # Stop the webhook dispatcher (and close its HTTP client) on shutdown, then defer to
        # NAT's own session-manager cleanup.
        if self._webhook_dispatcher is not None:
            await self._webhook_dispatcher.stop()
            self._webhook_dispatcher = None
        if self._progress_store is not None:
            await self._progress_store.close()
            self._progress_store = None
        await super().cleanup_session_managers()

    def build_app(self) -> FastAPI:
        app = super().build_app()
        config = AuthConfig.from_env()
        if config.enabled:
            app.add_middleware(JwtAuthMiddleware, config=config)
        else:
            logger.info("Inbound JWT auth is disabled (set VULN_ANALYSIS_AUTH_ENABLED=true to enable).")
        return app