FabledScribe/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.briefing import briefing_bp
from fabledassistant.routes.chat import chat_bp
from fabledassistant.routes.export import export_bp
from fabledassistant.routes.notes import notes_bp
from fabledassistant.routes.images import images_bp
from fabledassistant.routes.milestones import milestones_bp
from fabledassistant.routes.task_logs import task_logs_bp
from fabledassistant.routes.projects import projects_bp
from fabledassistant.routes.push import push_bp
from fabledassistant.routes.fable_mcp_dist import fable_mcp_dist_bp
from fabledassistant.routes.quick_capture import quick_capture_bp
from fabledassistant.routes.settings import settings_bp
from fabledassistant.routes.tasks import tasks_bp
from fabledassistant.routes.groups import groups_bp
from fabledassistant.routes.shares import shares_bp
from fabledassistant.routes.in_app_notifications import notifications_bp
from fabledassistant.routes.users import users_bp
from fabledassistant.routes.api_keys import api_keys_bp
from fabledassistant.routes.events import events_bp
from fabledassistant.routes.search import search_bp
from fabledassistant.routes.voice import voice_bp
from fabledassistant.routes.profile import profile_bp
from fabledassistant.routes.knowledge import knowledge_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),
    )

    # Validate config early so misconfigurations surface immediately with clear messages
    try:
        Config.validate()
    except ValueError as exc:
        logging.getLogger(__name__).error("Invalid configuration:\n%s", exc)
        raise

    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."
        )

    if not Config.TRUST_PROXY_HEADERS:
        logger.warning(
            "TRUST_PROXY_HEADERS is not set. If this instance is behind a reverse proxy "
            "(nginx, Caddy, Traefik) set TRUST_PROXY_HEADERS=true so rate limiting uses "
            "real client IPs rather than the proxy IP."
        )

    app.register_blueprint(admin_bp)
    app.register_blueprint(api)
    app.register_blueprint(auth_bp)
    app.register_blueprint(briefing_bp)
    app.register_blueprint(chat_bp)
    app.register_blueprint(export_bp)
    app.register_blueprint(images_bp)
    app.register_blueprint(milestones_bp)
    app.register_blueprint(notes_bp)
    app.register_blueprint(fable_mcp_dist_bp)
    app.register_blueprint(projects_bp)
    app.register_blueprint(push_bp)
    app.register_blueprint(quick_capture_bp)
    app.register_blueprint(settings_bp)
    app.register_blueprint(task_logs_bp)
    app.register_blueprint(tasks_bp)
    app.register_blueprint(groups_bp)
    app.register_blueprint(shares_bp)
    app.register_blueprint(notifications_bp)
    app.register_blueprint(users_bp)
    app.register_blueprint(api_keys_bp)
    app.register_blueprint(events_bp)
    app.register_blueprint(search_bp)
    app.register_blueprint(voice_bp)
    app.register_blueprint(profile_bp)
    app.register_blueprint(knowledge_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)
        # Downgrade noisy high-frequency / static paths to DEBUG
        _quiet = (
            request.path in {"/api/health", "/api/chat/status"}
            or request.path.startswith(("/static/", "/assets/", "/sw.js", "/manifest.json"))
        )
        log_fn = logger.debug if _quiet else logger.info
        log_fn(
            "%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'; worker-src '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
        from fabledassistant.services.push import ensure_vapid_keys

        start_cleanup_loop()
        start_log_retention_loop()
        start_notification_loop()
        ensure_vapid_keys()

        async def _warm_model(model: str) -> None:
            """Warm an already-installed model into VRAM (no pull)."""
            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": "2h"},
                    )
                logger.info("Warmed model '%s' into VRAM", model)
            except Exception:
                logger.warning("Failed to warm model '%s'", model, exc_info=True)

        async def _prime_kv_cache(user_id: int, model: str) -> None:
            """Send a minimal chat request to prime Ollama's KV cache with the user's system prompt.

            This ensures the next real user message only needs to process its own tokens
            rather than the full ~4,650-token system prompt, cutting TTFT from ~25s to <1s.
            The num_ctx must match what real requests will use so Ollama doesn't reload.
            """
            try:
                from fabledassistant.services.llm import build_context, pick_num_ctx
                messages, _ = await build_context(
                    user_id=user_id,
                    history=[],
                    current_note_id=None,
                    user_message=" ",
                )
                num_ctx = pick_num_ctx(messages)
                async with httpx.AsyncClient(timeout=120.0) as client:
                    await client.post(
                        f"{Config.OLLAMA_URL}/api/chat",
                        json={
                            "model": model,
                            "messages": messages,
                            "stream": False,
                            "options": {"num_predict": 1, "num_ctx": num_ctx},
                            "keep_alive": "2h",
                        },
                    )
                logger.info("Primed KV cache for user %d with model '%s' (num_ctx=%d)", user_id, model, num_ctx)
            except Exception:
                logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)

        async def _warm_user_models() -> None:
            """
            Warm whichever chat model(s) users have selected in Settings, then prime
            the KV cache with each user's system prompt so the first real message is fast.

            Only warms models that are already installed in Ollama — never auto-pulls.
            Falls back silently if no user preferences exist or Ollama is unreachable.
            """
            from sqlalchemy import select as sa_select

            from fabledassistant.models import async_session
            from fabledassistant.models.setting import Setting

            # 1. Collect (user_id, model) pairs for all users with a saved default_model.
            try:
                async with async_session() as session:
                    rows = await session.execute(
                        sa_select(Setting.user_id, Setting.value).where(
                            Setting.key == "default_model",
                            Setting.value.isnot(None),
                            Setting.value != "",
                        )
                    )
                    user_model_pairs: list[tuple[int, str]] = list(rows)
            except Exception:
                logger.debug("Could not read user model preferences from DB", exc_info=True)
                return

            if not user_model_pairs:
                logger.debug("No user model preferences found; skipping warm-up")
                return

            # 2. Ask Ollama which models are currently installed.
            try:
                async with httpx.AsyncClient(timeout=10.0) as client:
                    resp = await client.get(f"{Config.OLLAMA_URL}/api/tags")
                    resp.raise_for_status()
                    installed: set[str] = {m["name"] for m in resp.json().get("models", [])}
            except Exception:
                logger.debug("Could not reach Ollama to check installed models", exc_info=True)
                return

            # 3. Warm each unique model, then prime KV cache per user.
            warmed: set[str] = set()
            for user_id_val, model in user_model_pairs:
                base = model.removesuffix(":latest")
                if model in installed or f"{base}:latest" in installed or base in installed:
                    if model not in warmed:
                        await _warm_model(model)
                        warmed.add(model)
                    await _prime_kv_cache(user_id_val, model)
                else:
                    logger.info(
                        "User-preferred model '%s' is not installed; skipping warm-up "
                        "(install it via Settings → Models to enable auto-warm)",
                        model,
                    )

        async def _pull_model(model: str, warm: bool = False) -> None:
            try:
                await ensure_model(model)
            except Exception:
                logger.warning(
                    "Failed to ensure model '%s'",
                    model,
                    exc_info=True,
                )
                return
            if warm:
                await _warm_model(model)

        # Warm user-preferred chat models that are already installed.
        # Also ensure the embedding model is pulled (no warm needed).
        asyncio.create_task(_warm_user_models())
        asyncio.create_task(_pull_model(Config.EMBEDDING_MODEL, warm=False))
        asyncio.create_task(_pull_model(Config.OLLAMA_BACKGROUND_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)
            try:
                from fabledassistant.services.projects import backfill_project_summaries
                await backfill_project_summaries()
            except Exception:
                logger.warning("Project summary backfill failed", exc_info=True)
            try:
                from fabledassistant.services.embeddings import backfill_rss_item_embeddings
                await backfill_rss_item_embeddings()
            except Exception:
                logger.warning("RSS embedding backfill failed", exc_info=True)
            try:
                from fabledassistant.services.embeddings import backfill_rss_article_content
                await backfill_rss_article_content()
            except Exception:
                logger.warning("RSS article content backfill failed", exc_info=True)

        asyncio.create_task(_delayed_backfill())

        # Start briefing scheduler
        from fabledassistant.services.briefing_scheduler import start_briefing_scheduler
        await start_briefing_scheduler(asyncio.get_running_loop())

        # Voice model loading (enabled via Admin → Config in the UI, or VOICE_ENABLED env var)
        from fabledassistant.services.stt import load_stt_model
        from fabledassistant.services.tts import load_tts_model
        asyncio.create_task(load_stt_model())
        asyncio.create_task(load_tts_model())

    @app.after_serving
    async def shutdown():
        from fabledassistant.services.briefing_scheduler import stop_briefing_scheduler
        stop_briefing_scheduler()

    @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

    # File extensions that belong to the SPA or its assets.
    # Anything else with an extension is not a valid app path and gets a hard 404.
    _SPA_EXTENSIONS = {
        "", ".html", ".js", ".css", ".ico", ".png", ".jpg", ".jpeg",
        ".svg", ".webp", ".woff", ".woff2", ".ttf", ".otf", ".map", ".json",
    }

    @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 a real static file first
        path = request.path.lstrip("/")
        file_path = STATIC_DIR / path
        if path and file_path.is_file():
            return await send_from_directory(STATIC_DIR, path)

        # Reject paths with file extensions the app doesn't serve.
        # This turns scanner probes (.php, .asp, .cgi, etc.) into honest 404s
        # instead of serving them the Vue SPA with a misleading 200.
        from pathlib import PurePosixPath
        suffix = PurePosixPath(request.path).suffix.lower()
        if suffix not in _SPA_EXTENSIONS:
            return "", 404

        # SPA fallback for clean client-side routes (/notes/123, /chat, etc.)
        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