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.chat import chat_bp
from fabledassistant.routes.journal import journal_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(chat_bp)
    app.register_blueprint(journal_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' 'wasm-unsafe-eval'; "
            "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; "
            "img-src 'self' data: blob:; "
            "connect-src 'self'; "
            "font-src 'self' data: https://fonts.gstatic.com; "
            "object-src 'none'; "
            "base-uri 'self'; worker-src 'self' blob:;"
        )

        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)."""
            from fabledassistant.services.llm import keep_alive_for
            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": keep_alive_for(model)},
                    )
                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, keep_alive_for, pick_num_ctx
                from fabledassistant.services.tools import get_tools_for_user
                messages, _ = await build_context(
                    user_id=user_id,
                    history=[],
                    current_note_id=None,
                    user_message=" ",
                )
                # Include tool schemas so num_ctx matches real chat requests.
                tools = await get_tools_for_user(user_id)
                num_ctx = pick_num_ctx(messages, tools=tools)
                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": keep_alive_for(model),
                        },
                    )
                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:
            """
            Pull any user-configured models that are missing from Ollama, then warm
            them and prime the KV cache with each user's system prompt.

            Handles both default_model (chat) and background_model user overrides.
            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 all user model preferences (both chat and background).
            try:
                async with async_session() as session:
                    rows = await session.execute(
                        sa_select(Setting.user_id, Setting.key, Setting.value).where(
                            Setting.key.in_(["default_model", "background_model"]),
                            Setting.value.isnot(None),
                            Setting.value != "",
                        )
                    )
                    settings_rows: list[tuple[int, str, str]] = list(rows)
            except Exception:
                logger.debug("Could not read user model preferences from DB", exc_info=True)
                return

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

            # 2. Build the set of unique models to ensure, and the list of
            #    (user_id, chat_model) pairs for KV-cache priming.
            all_models: set[str] = set()
            user_chat_models: list[tuple[int, str]] = []
            for user_id_val, key, model in settings_rows:
                all_models.add(model)
                if key == "default_model":
                    user_chat_models.append((user_id_val, model))

            # 3. 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()
                    raw_installed: set[str] = {m["name"] for m in resp.json().get("models", [])}
                    installed: set[str] = raw_installed | {
                        n.removesuffix(":latest") for n in raw_installed if n.endswith(":latest")
                    }
            except Exception:
                logger.debug("Could not reach Ollama to check installed models", exc_info=True)
                return

            # 4. Pull any user-configured models that are missing.
            for model in all_models:
                if model not in installed:
                    logger.info("User-configured model '%s' not installed; pulling...", model)
                    await _pull_model(model)
                    installed.add(model)

            # 5. Warm each unique chat model, then prime KV cache per user.
            warmed: set[str] = set()
            for user_id_val, model in user_chat_models:
                if model in installed:
                    if model not in warmed:
                        await _warm_model(model)
                        warmed.add(model)
                    await _prime_kv_cache(user_id_val, 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)

        # Pull supporting models without warming them — embedding model and
        # the configured background model load on demand without competing
        # for VRAM. The chat model is warmed by _warm_user_models() which
        # uses each user's *actual* default_model setting; warming
        # Config.OLLAMA_MODEL unconditionally (the OLD behaviour) blasted
        # the system default into VRAM ahead of the user's real preference,
        # which on a single-GPU setup pushed the user's chat model out
        # before they ever sent a message.
        asyncio.create_task(_pull_model(Config.OLLAMA_MODEL, warm=False))
        asyncio.create_task(_pull_model(Config.EMBEDDING_MODEL, warm=False))
        asyncio.create_task(_pull_model(Config.OLLAMA_BACKGROUND_MODEL, warm=False))
        asyncio.create_task(_warm_user_models())

        # 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)

        asyncio.create_task(_delayed_backfill())

        # Start journal scheduler (per-user daily prep generation)
        from fabledassistant.services.journal_scheduler import start_journal_scheduler
        start_journal_scheduler(asyncio.get_running_loop())

        # Start event scheduler (reminders + CalDAV pull sync)
        from fabledassistant.services.event_scheduler import start_event_scheduler
        start_event_scheduler(asyncio.get_running_loop())

        # Start version-pinning scheduler (daily auto-pin scan at 03:00 UTC)
        from fabledassistant.services.version_pinning_scheduler import (
            start_version_pinning_scheduler,
        )
        start_version_pinning_scheduler(asyncio.get_running_loop())

        # Start curator scheduler (15-min sweep of journal conversations)
        from fabledassistant.services.curator_scheduler import start_curator_scheduler
        start_curator_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.journal_scheduler import stop_journal_scheduler
        stop_journal_scheduler()
        from fabledassistant.services.event_scheduler import stop_event_scheduler
        stop_event_scheduler()
        from fabledassistant.services.version_pinning_scheduler import (
            stop_version_pinning_scheduler,
        )
        stop_version_pinning_scheduler()
        from fabledassistant.services.curator_scheduler import stop_curator_scheduler
        stop_curator_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