FabledGrimoire/src/fabledassistant/app.py

import logging
import time
import traceback as tb_module
from pathlib import Path

import httpx

from quart import Quart, g, jsonify, make_response, request, send_from_directory

from fabledassistant.config import Config
from fabledassistant.routes.admin import admin_bp
from fabledassistant.routes.api import api
from fabledassistant.routes.auth import auth_bp
from fabledassistant.routes.chat import chat_bp
from fabledassistant.routes.notes import notes_bp
from fabledassistant.routes.settings import settings_bp
from fabledassistant.routes.tasks import tasks_bp

STATIC_DIR = Path(__file__).parent / "static"
logger = logging.getLogger(__name__)


def create_app() -> Quart:
    # Configure logging
    logging.basicConfig(
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        level=getattr(logging, Config.LOG_LEVEL.upper(), logging.INFO),
    )

    app = Quart(__name__, static_folder=None)
    app.secret_key = Config.SECRET_KEY
    app.config["SESSION_COOKIE_HTTPONLY"] = True
    app.config["SESSION_COOKIE_SAMESITE"] = "Lax"
    app.config["SESSION_COOKIE_SECURE"] = Config.SECURE_COOKIES

    if Config.SECRET_KEY == "dev-secret-change-me":
        logger.warning(
            "SECRET_KEY is set to the default value — session cookies are insecure. "
            "Set SECRET_KEY or SECRET_KEY_FILE for production use."
        )

    app.register_blueprint(admin_bp)
    app.register_blueprint(api)
    app.register_blueprint(auth_bp)
    app.register_blueprint(chat_bp)
    app.register_blueprint(notes_bp)
    app.register_blueprint(settings_bp)
    app.register_blueprint(tasks_bp)

    @app.before_request
    async def before_request():
        g.request_start = time.monotonic()

    @app.after_request
    async def after_request(response):
        duration = time.monotonic() - getattr(g, "request_start", time.monotonic())
        duration_ms = round(duration * 1000, 1)
        logger.info(
            "%s %s %s %.1fms",
            request.method,
            request.path,
            response.status_code,
            duration_ms,
        )

        # Log usage for API requests (skip logs endpoint to avoid recursion)
        if request.path.startswith("/api/") and not request.path.startswith("/api/admin/logs"):
            try:
                from fabledassistant.services.logging import log_usage

                user = getattr(g, "user", None)
                await log_usage(
                    user_id=user.id if user else None,
                    username=user.username if user else None,
                    endpoint=request.path,
                    method=request.method,
                    status_code=response.status_code,
                    duration_ms=duration_ms,
                )
            except Exception:
                logger.debug("Failed to log usage", exc_info=True)

        response.headers.setdefault("X-Content-Type-Options", "nosniff")
        response.headers.setdefault("X-Frame-Options", "DENY")
        response.headers.setdefault("Referrer-Policy", "strict-origin-when-cross-origin")
        response.headers.setdefault(
            "Content-Security-Policy",
            "default-src 'self'; script-src 'self' 'unsafe-inline'; "
            "style-src 'self' 'unsafe-inline'; img-src 'self' data: blob:; "
            "connect-src 'self'; font-src 'self' data:; object-src 'none'; base-uri 'self';"
        )

        return response

    @app.before_serving
    async def startup():
        import asyncio

        from fabledassistant.services.embeddings import backfill_note_embeddings
        from fabledassistant.services.generation_buffer import start_cleanup_loop
        from fabledassistant.services.llm import ensure_model
        from fabledassistant.services.logging import start_log_retention_loop
        from fabledassistant.services.notifications import start_notification_loop

        start_cleanup_loop()
        start_log_retention_loop()
        start_notification_loop()

        async def _pull_model(model: str, warm: bool = False) -> None:
            try:
                await ensure_model(model)
            except Exception:
                logger.warning(
                    "Failed to ensure model '%s' — chat may not work until model is available",
                    model,
                    exc_info=True,
                )
                return
            if warm:
                try:
                    async with httpx.AsyncClient(timeout=300.0) as client:
                        await client.post(
                            f"{Config.OLLAMA_URL}/api/generate",
                            json={"model": model, "prompt": "", "keep_alive": "30m"},
                        )
                    logger.info("Warmed model '%s' into VRAM", model)
                except Exception:
                    logger.warning("Failed to warm model '%s'", model, exc_info=True)

        # Pull main model and (if configured) a separate intent model concurrently,
        # then warm them into VRAM so the first user request doesn't cold-load both
        # simultaneously (which causes Ollama 500 races).
        # Fire-and-forget so pulls/warming don't block startup.
        models_to_warm = {Config.OLLAMA_MODEL}
        if Config.OLLAMA_INTENT_MODEL and Config.OLLAMA_INTENT_MODEL != Config.OLLAMA_MODEL:
            models_to_warm.add(Config.OLLAMA_INTENT_MODEL)
        for _model in models_to_warm:
            asyncio.create_task(_pull_model(_model, warm=True))
        # Also pull the embedding model (nomic-embed-text by default), but no need to warm it.
        if Config.EMBEDDING_MODEL not in models_to_warm:
            asyncio.create_task(_pull_model(Config.EMBEDDING_MODEL, warm=False))

        # After models are pulled, backfill embeddings for existing notes.
        # Runs in the background so it never blocks the server from accepting requests.
        async def _delayed_backfill() -> None:
            await asyncio.sleep(30)  # Give Ollama time to load the embedding model
            try:
                await backfill_note_embeddings()
            except Exception:
                logger.warning("Embedding backfill failed", exc_info=True)

        asyncio.create_task(_delayed_backfill())

    @app.route("/")
    async def serve_index():
        resp = await make_response(
            await send_from_directory(STATIC_DIR, "index.html")
        )
        resp.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
        return resp

    @app.errorhandler(404)
    async def handle_404(error):
        # Return JSON 404 for API routes
        if request.path.startswith("/api/"):
            return jsonify({"error": "Not found"}), 404
        # Try to serve static file
        path = request.path.lstrip("/")
        file_path = STATIC_DIR / path
        if path and file_path.is_file():
            return await send_from_directory(STATIC_DIR, path)
        # SPA fallback
        resp = await make_response(
            await send_from_directory(STATIC_DIR, "index.html")
        )
        resp.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
        return resp

    @app.errorhandler(500)
    async def handle_500(error):
        logger.exception("Internal server error on %s %s", request.method, request.path)

        try:
            from fabledassistant.services.logging import log_error

            user = getattr(g, "user", None)
            await log_error(
                user_id=user.id if user else None,
                username=user.username if user else None,
                endpoint=request.path,
                method=request.method,
                error_type=type(error).__name__,
                error_message=str(error),
                traceback=tb_module.format_exc(),
            )
        except Exception:
            logger.debug("Failed to log error", exc_info=True)

        if request.path.startswith("/api/"):
            return jsonify({"error": "Internal server error"}), 500
        return "Internal Server Error", 500

    return app