app/main.py

import os
from pathlib import Path
from contextlib import asynccontextmanager

from dotenv import load_dotenv

# Must run before any local import that reads env vars. Resolve explicitly so
# Electron, uv, and direct uvicorn launches all see backend/.env regardless of cwd.
load_dotenv(Path(__file__).resolve().parents[1] / ".env")

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from uvicorn.protocols.utils import ClientDisconnected
from fastapi.middleware.cors import CORSMiddleware

from app.routers import agent_control, annotations, citation_graph, control_room, health, library, projects, regions, sandbox, session, review, chat, draft, wiki, voice, visual_lessons
from app.websockets.handlers import get_connection_manager, get_db, handle_event


@asynccontextmanager
async def lifespan(app: FastAPI):
    import asyncio
    import logging
    import os
    import time
    logger = logging.getLogger(__name__)

    # Observability boundary: a sibling startup step to the Cognee bootstrap
    # below (same timed pattern). Best-effort -> a telemetry init failure must
    # never block the app from serving, so it degrades to disabled internally.
    from app.observability.bootstrap import initialize_observability, shutdown_observability
    from app.observability.config import get_observability_config

    obs_cfg = get_observability_config()
    try:
        obs_t0 = time.perf_counter()
        initialize_observability(app=app, config=obs_cfg)
        if obs_cfg.mode != "disabled":
            logger.warning(
                "Observability init (mode=%s) took %.2fs",
                obs_cfg.mode,
                time.perf_counter() - obs_t0,
            )
    except Exception:
        logger.exception("Observability init failed")

    async def _shutdown_observability() -> None:
        # Hard wall-clock deadline: a hung/unreachable OTLP collector must not
        # stall process teardown on flush.
        try:
            await asyncio.wait_for(asyncio.to_thread(shutdown_observability, 5.0), timeout=6.0)
        except Exception:
            logger.exception("Observability shutdown failed")

    try:
        if os.getenv("DEPLOYMENT_ENV", "desktop") == "demo":
            print("Demo mode: skipping Cognee bootstrap")
            yield
            return

        await _run_cognee_bootstrap(logger)
        yield
    finally:
        await _shutdown_observability()


async def _run_cognee_bootstrap(logger) -> None:
    import asyncio
    import logging
    import os
    import time
    from pathlib import Path

    # Cerebras has no embeddings endpoint -> Cognee's default embedding model
    # (openai/text-embedding-3-large) would otherwise be routed through the
    # OPENAI_API_BASE override above and 404. Use a local, fully offline
    # embedding model instead (matches CLAUDE.md: Cognee never writes to cloud).
    os.environ.setdefault("EMBEDDING_PROVIDER", "fastembed")
    os.environ.setdefault("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
    os.environ.setdefault("EMBEDDING_DIMENSIONS", "384")
    skill_root = str(Path(__file__).resolve().parent / "memory_skills")
    existing_skill_roots = os.environ.get("COGNEE_SKILL_SOURCE_ROOTS", "")
    if skill_root not in existing_skill_roots.split(os.pathsep):
        os.environ["COGNEE_SKILL_SOURCE_ROOTS"] = (
            skill_root if not existing_skill_roots else existing_skill_roots + os.pathsep + skill_root
        )
    
    # Disable multi-tenant access control for Cognee 1.2+ since StudyBuddy uses a single local user
    os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "false"

    t0 = time.perf_counter()
    import cognee
    try:
        import importlib
        from app.services.cognee_bootstrap import configure_cognee_llm

        cognee_llm_client = importlib.import_module(
            "cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client"
        )

        configure_cognee_llm(
            cognee.config,
            clear_llm_client_cache=cognee_llm_client._get_llm_client_cached.cache_clear,
        )
        # Suppress expected "No data found" and "DatabaseNotCreatedError" logs on fresh installs
        logging.getLogger("cognee.shared.logging_utils").setLevel(logging.CRITICAL)
        
        root = str(Path.home() / ".studybuddy" / "cognee")
        cognee.config.data_root_directory(root)
        cognee.config.system_root_directory(f"{root}/system")
        from cognee.infrastructure.databases.relational.create_db_and_tables import create_db_and_tables
        await create_db_and_tables()
        logger.warning("Cognee mandatory bootstrap took %.2fs", time.perf_counter() - t0)

        # Bootstrap the one cross-project student-profile dataset. Project
        # research memory is intentionally Chroma-backed and never creates a
        # Cognee dataset.
        from app.services.student_memory import StudentMemoryService
        async def _warm_profile_dataset() -> None:
            warm_t0 = time.perf_counter()
            try:
                await StudentMemoryService().ensure_profile_dataset()
                logger.warning("Cognee profile dataset warmup took %.2fs", time.perf_counter() - warm_t0)
            except Exception:
                logger.exception("Cognee profile dataset warmup failed")
        asyncio.create_task(_warm_profile_dataset())
    except Exception as e:
        print("Cognee setup error:", e)


app = FastAPI(title="ResearchMate API", lifespan=lifespan)

_origins_env = os.getenv("ALLOWED_ORIGINS", "")
_origins = (
    _origins_env.split(",")
    if _origins_env
    else [
        "http://localhost:5173",
        "http://127.0.0.1:5173",
        "http://localhost:5174",
        "http://127.0.0.1:5174",
    ]
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(health.router)
app.include_router(control_room.router)
app.include_router(agent_control.router)
app.include_router(library.router)
app.include_router(projects.router)
app.include_router(sandbox.router)
app.include_router(session.router)
app.include_router(annotations.router)
app.include_router(regions.router)
app.include_router(review.router)
app.include_router(citation_graph.router)
app.include_router(chat.router, prefix="/chat", tags=["Chat"])
app.include_router(draft.router)
app.include_router(wiki.router)
app.include_router(voice.router)
app.include_router(visual_lessons.router)


@app.websocket("/ws/{project_id}")
async def websocket_endpoint(ws: WebSocket, project_id: str):
    import asyncio
    import logging
    logger = logging.getLogger(__name__)
    cm = get_connection_manager()
    background_events: set[asyncio.Task[None]] = set()

    async def run_event(event_type: str, data: dict) -> None:
        try:
            await handle_event(project_id, event_type, data)
            if event_type in {"CHAT_TURN", "CONTEXT_CARD_REQUEST", "VISUALIZATION_REQUEST", "CHAT_VISUAL_REQUEST", "CHAT_COMPACT"}:
                logger.info("websocket event done project=%s type=%s", project_id, event_type)
        except (WebSocketDisconnect, ClientDisconnected, RuntimeError) as exc:
            logger.info("websocket event aborted project=%s type=%s detail=%s", project_id, event_type, exc)
        except Exception:
            logger.exception("handle_event failed for session %s event %s", project_id, event_type)
            await cm.send(project_id, "ERROR", {
                "event_type": event_type,
                "message": "Something went wrong processing that -> please try again.",
            })

    await cm.connect(project_id, ws)
    logger.info("websocket connected project=%s", project_id)
    try:
        while True:
            msg = await ws.receive_json()
            event_type = msg.get("type", "")
            if event_type in {"CHAT_TURN", "CONTEXT_CARD_REQUEST", "VISUALIZATION_REQUEST", "CHAT_VISUAL_REQUEST"}:
                data = msg.get("data", {}) or {}
                logger.info(
                    "websocket event start project=%s type=%s selection_len=%s has_image=%s",
                    project_id,
                    event_type,
                    len(data.get("selection_text") or ""),
                    bool(data.get("selection_image_base64")),
                )
            data = msg.get("data", {}) or {}
            if event_type in {"CHAT_TURN", "CONTEXT_CARD_REQUEST", "VISUALIZATION_REQUEST", "CHAT_VISUAL_REQUEST", "CHAT_COMPACT"}:
                task = asyncio.create_task(run_event(event_type, data))
                background_events.add(task)
                task.add_done_callback(background_events.discard)
            else:
                await run_event(event_type, data)
    except WebSocketDisconnect as exc:
        logger.info("websocket disconnected project=%s code=%s", project_id, getattr(exc, "code", None))
        cm.disconnect(project_id, ws)
    except RuntimeError as exc:
        logger.info("websocket closed project=%s detail=%s", project_id, exc)
        cm.disconnect(project_id, ws)
    finally:
        for task in background_events:
            task.cancel()
        if background_events:
            await asyncio.gather(*background_events, return_exceptions=True)