Skyvern/skyvern/forge/api_app.py

"""FastAPI application factory for server installs."""

# The server-extra guard must run before FastAPI/Starlette imports.
# ruff: noqa: E402

import os
import uuid
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Any, AsyncGenerator, Awaitable, Callable

import structlog
from pydantic import ValidationError

from skyvern.exceptions import require_server_extra_modules

require_server_extra_modules("skyvern.forge.api_app", ("fastapi", "starlette", "starlette_context"))

from fastapi import FastAPI, Response, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.openapi.utils import get_openapi
from fastapi.responses import JSONResponse
from starlette.requests import HTTPConnection, Request
from starlette_context.middleware import RawContextMiddleware
from starlette_context.plugins.base import Plugin

from skyvern.config import _ensure_sqlite_dir, settings
from skyvern.exceptions import SkyvernHTTPException
from skyvern.forge import app as forge_app
from skyvern.forge.forge_app_initializer import start_forge_app
from skyvern.forge.request_logging import log_raw_request_middleware
from skyvern.forge.sdk.core import skyvern_context
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
from skyvern.forge.sdk.db.exceptions import NotFoundError
from skyvern.forge.sdk.db.models import Base
from skyvern.forge.sdk.routes import internal_auth
from skyvern.forge.sdk.routes.google_oauth import google_oauth_router
from skyvern.forge.sdk.routes.google_sheets import google_sheets_router
from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router, legacy_v2_router
from skyvern.forge.sdk.services.local_org_auth_token_service import (
    ensure_local_api_key,
    ensure_local_org,
    fingerprint_token,
    regenerate_local_api_key,
)
from skyvern.services.cleanup_service import start_cleanup_scheduler, stop_cleanup_scheduler

LOG = structlog.get_logger()


def format_validation_errors(exc: ValidationError) -> str:
    """Format a Pydantic ValidationError into a human-readable string.

    Filters out uninformative path segments ('__root__', 'body') and joins
    multiple errors with '; '.
    """
    error_messages = []
    for error in exc.errors():
        loc = " -> ".join(str(part) for part in error["loc"] if part not in ("__root__", "body"))
        msg = error["msg"]
        if loc:
            error_messages.append(f"{loc}: {msg}")
        else:
            error_messages.append(msg)
    return (
        "; ".join(error_messages)
        if error_messages
        else "A validation error occurred. Please check your input and try again."
    )


class ExecutionDatePlugin(Plugin):
    key = "execution_date"

    async def process_request(self, request: Request | HTTPConnection) -> datetime:
        return datetime.now()


def custom_openapi(app: FastAPI) -> dict:
    if app.openapi_schema:
        return app.openapi_schema
    openapi_schema = get_openapi(
        title="Skyvern API",
        version="1.0.0",
        description="API for Skyvern",
        routes=app.routes,
    )
    openapi_schema["servers"] = [
        {"url": "https://api.skyvern.com", "x-fern-server-name": "Cloud"},
        {"url": "https://api-staging.skyvern.com", "x-fern-server-name": "Staging"},
        {"url": "http://localhost:8000", "x-fern-server-name": "Local"},
    ]
    app.openapi_schema = openapi_schema
    return app.openapi_schema


async def _bootstrap_sqlite() -> None:
    """Auto-bootstrap SQLite on first server start.

    Creates tables, a local org, and an API key so that
    ``skyvern run server`` works out of the box with zero configuration.
    Idempotent: skips if the org already exists.
    """
    _ensure_sqlite_dir(settings.DATABASE_STRING)

    db = forge_app.DATABASE
    async with db.engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    # Preserve an existing API key if it's a real value (not the skeleton default).
    # settings.SKYVERN_API_KEY already incorporates env vars and .env via pydantic-settings.
    existing_key = settings.SKYVERN_API_KEY if settings.SKYVERN_API_KEY != "PLACEHOLDER" else None

    if existing_key:
        preserved = await ensure_local_api_key(existing_key)
        if preserved is not None:
            api_key, org_id = preserved
            LOG.info(
                "Existing SKYVERN_API_KEY detected — preserving env value and syncing it into the local SQLite DB.",
                organization_id=org_id,
                api_key_fingerprint=fingerprint_token(api_key),
            )
            return

        LOG.warning(
            "Existing SKYVERN_API_KEY could not be preserved for local SQLite bootstrap; generating a new local key.",
        )

    organization = await ensure_local_org()
    existing_token = await db.organizations.get_valid_org_auth_token(organization.organization_id, "api")
    if existing_token is not None:
        LOG.info("SQLite database already bootstrapped", organization_id=organization.organization_id)
        return

    api_key, org_id, backend_env, frontend_env = await regenerate_local_api_key()
    LOG.info(
        "SQLite bootstrap complete — local org and API key created",
        organization_id=org_id,
        api_key_fingerprint=fingerprint_token(api_key),
        env_file_written=backend_env,
    )


@asynccontextmanager
async def lifespan(fastapi_app: FastAPI) -> AsyncGenerator[None, Any]:
    """Lifespan context manager for FastAPI app startup and shutdown."""

    LOG.info("Server started")

    # Auto-bootstrap SQLite database on first server start.
    # Re-raise on failure — a server with no tables/org/API key is
    # useless and would produce confusing 401s on every request.
    if settings.is_sqlite():
        await _bootstrap_sqlite()

    if forge_app.api_app_startup_event:
        LOG.info("Calling api app startup event")
        try:
            await forge_app.api_app_startup_event(fastapi_app)
        except Exception:
            LOG.exception("Failed to execute api app startup event")

    # Close browser sessions left active by a previous process
    try:
        await forge_app.PERSISTENT_SESSIONS_MANAGER.cleanup_stale_sessions()
    except Exception:
        LOG.exception("Failed to clean up stale browser sessions")

    # Start cleanup scheduler if enabled
    cleanup_task = start_cleanup_scheduler()
    if cleanup_task:
        LOG.info("Cleanup scheduler started")

    # Start MCP sub-application lifespan if mounted. Starlette Mount does NOT
    # forward lifespan events to sub-apps, so we must enter the MCP app's
    # lifespan here. This initializes the streamable-http session manager's
    # task group which is required for handling MCP requests.
    mcp_app = getattr(fastapi_app.state, "mcp_starlette_app", None)
    if mcp_app:
        async with mcp_app.lifespan(mcp_app):
            LOG.info("MCP remote server lifespan started")
            yield
        LOG.info("MCP remote server lifespan stopped")
    else:
        yield

    # Stop cleanup scheduler
    await stop_cleanup_scheduler()

    if forge_app.api_app_shutdown_event:
        LOG.info("Calling api app shutdown event")
        try:
            await forge_app.api_app_shutdown_event()
        except Exception:
            LOG.exception("Failed to execute api app shutdown event")
    LOG.info("Server shutting down")


def create_api_app() -> FastAPI:
    """
    Start the agent server.
    """
    # CRITICAL: Initialize OTEL FIRST, before any other code runs.
    # This must happen before start_forge_app() because that function
    # creates database connections. If we don't instrument the libraries
    # first, the DB spans won't be children of the HTTP request spans.
    # NOTE: The cloud import is lazy (not at module level) to avoid triggering
    # cloud/__init__.py side effects when skyvern is used in embedded mode.
    otel_setup_cls: Any = None
    if settings.OTEL_ENABLED:
        try:
            from cloud.observability.otel_setup import OTELSetup  # noqa: PLC0415

            otel_setup_cls = OTELSetup
        except ImportError:
            pass
    if otel_setup_cls is not None:
        try:
            otel = otel_setup_cls.get_instance()
            otel.initialize_tracer_provider()
            LOG.info("OTEL tracer provider initialized before forge app creation")
        except Exception as e:
            LOG.warning("Failed to initialize OTEL tracer provider early", error=str(e))

    forge_app_instance = start_forge_app()

    # Initialize Laminar tracing after ForgeApp so auto-instrumentation works.
    lmnr_api_key = os.environ.get("LMNR_PROJECT_API_KEY")
    if lmnr_api_key:
        try:
            from lmnr import Laminar  # noqa: PLC0415

            lmnr_base_url = os.environ.get("LMNR_BASE_URL", "http://localhost")
            lmnr_grpc_port = int(os.environ.get("LMNR_GRPC_PORT", "8001"))
            lmnr_http_port = int(os.environ.get("LMNR_HTTP_PORT", "8000"))
            Laminar.initialize(
                project_api_key=lmnr_api_key,
                base_url=lmnr_base_url,
                grpc_port=lmnr_grpc_port,
                http_port=lmnr_http_port,
            )
            LOG.info("Laminar tracing initialized", base_url=lmnr_base_url, grpc_port=lmnr_grpc_port)
        except Exception as e:
            LOG.warning("Failed to initialize Laminar tracing", error=str(e))

    fastapi_app = FastAPI(lifespan=lifespan)

    # Add CORS middleware
    fastapi_app.add_middleware(
        CORSMiddleware,
        allow_origins=settings.ALLOWED_ORIGINS,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    fastapi_app.include_router(base_router, prefix="/v1")
    fastapi_app.include_router(legacy_base_router, prefix="/api/v1")
    fastapi_app.include_router(legacy_v2_router, prefix="/api/v2")
    # Google OAuth is mounted on both the canonical ``/v1`` (matches base_router)
    # and the legacy ``/api/v1`` (matches legacy_base_router) — same dual-prefix
    # pattern as internal_auth below. Both registrations are intentional; do not
    # remove either as a "duplicate". ``include_in_schema=False`` keeps the OAuth
    # endpoints out of the public OpenAPI/Swagger surface — they're consumed by
    # the frontend, not by SDK users.
    fastapi_app.include_router(google_oauth_router, prefix="/v1/google", include_in_schema=False)
    fastapi_app.include_router(google_oauth_router, prefix="/api/v1/google", include_in_schema=False)
    fastapi_app.include_router(google_sheets_router, prefix="/v1/google/sheets", include_in_schema=False)
    fastapi_app.include_router(google_sheets_router, prefix="/api/v1/google/sheets", include_in_schema=False)

    # local dev endpoints
    if settings.ENV == "local":
        fastapi_app.include_router(internal_auth.router, prefix="/v1")
        fastapi_app.include_router(internal_auth.router, prefix="/api/v1")
        fastapi_app.include_router(internal_auth.router, prefix="/api/v2")

    fastapi_app.openapi = lambda: custom_openapi(fastapi_app)

    fastapi_app.add_middleware(
        RawContextMiddleware,
        plugins=(
            # TODO (suchintan): We should set these up
            ExecutionDatePlugin(),
            # RequestIdPlugin(),
            # UserAgentPlugin(),
        ),
    )

    @fastapi_app.exception_handler(NotFoundError)
    async def handle_not_found_error(request: Request, exc: NotFoundError) -> Response:
        return Response(status_code=status.HTTP_404_NOT_FOUND)

    @fastapi_app.exception_handler(SkyvernHTTPException)
    async def handle_skyvern_http_exception(request: Request, exc: SkyvernHTTPException) -> JSONResponse:
        return JSONResponse(status_code=exc.status_code, content={"detail": exc.message})

    @fastapi_app.exception_handler(ValidationError)
    async def handle_pydantic_validation_error(request: Request, exc: ValidationError) -> JSONResponse:
        detail = format_validation_errors(exc)
        return JSONResponse(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            content={"detail": detail},
        )

    @fastapi_app.exception_handler(Exception)
    async def unexpected_exception(request: Request, exc: Exception) -> JSONResponse:
        LOG.exception("Unexpected error in agent server.", exc_info=exc)
        return JSONResponse(status_code=500, content={"error": f"Unexpected error: {exc}"})

    @fastapi_app.middleware("http")
    async def request_middleware(request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        curr_ctx = skyvern_context.current()
        if not curr_ctx:
            request_id = str(uuid.uuid4())
            skyvern_context.set(SkyvernContext(request_id=request_id))
        elif not curr_ctx.request_id:
            curr_ctx.request_id = str(uuid.uuid4())

        try:
            return await call_next(request)
        finally:
            skyvern_context.reset()

    @fastapi_app.middleware("http")
    async def raw_request_logging(request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        return await log_raw_request_middleware(request, call_next)

    if forge_app_instance.setup_api_app:
        forge_app_instance.setup_api_app(fastapi_app)

    return fastapi_app