ss-tools/backend/src/app.py

# #region AppModule [C:5] [TYPE Module] [SEMANTICS fastapi, websocket, startup, scheduler, middleware]
# @defgroup Module Module group.
# @BRIEF The main entry point for the FastAPI application.
# @LAYER API
# @RELATION DEPENDS_ON -> [ApiRoutesModule]
# @INVARIANT All WebSocket connections must be properly cleaned up on disconnect.
# @INVARIANT All WebSocket connections must be authenticated via JWT or API key token (see [SEC:C-3]).
# @PRE Python environment and dependencies installed; configuration database available.
# @POST FastAPI app instance is created, middleware configured, and routes registered.
# @SIDE_EFFECT Starts background scheduler and binds network ports for HTTP/WS traffic.
# @DATA_CONTRACT [HTTP Request | WS Message] -> [HTTP Response | JSON Log Stream]
# @RATIONALE Duplicate @RELATION and @INVARIANT lines removed from header. Import sorting unified via ruff isort (I) rule across src/ — 139 fixes applied.
# @REJECTED Monolithic route registration in app.py was rejected — route groups are now modular
#           in separate router modules under api/routes/ to keep AppModule focused on middleware,
#           lifespan, and WebSocket endpoint wiring. Embedding lifespan logic directly in module
#           scope was rejected — the asynccontextmanager pattern ensures clean startup/shutdown
#           lifecycle without module-level race conditions.

import asyncio
from contextlib import asynccontextmanager
import os
from pathlib import Path
import sys

# project_root is used for static files mounting
project_root = Path(__file__).resolve().parent.parent.parent

from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.sessions import SessionMiddleware
from typing import Any

from .api import auth
from .api.routes import (
    admin,
    admin_api_keys,
    agent_conversations,
    agent_status,
    agent_superset,
    agent_superset_explore,
    assistant,
    clean_release,
    clean_release_v2,
    dashboards,
    dataset_review,
    datasets,
    encryption_health,
    environments,
    git,
    health,
    llm,
    maintenance,
    mappings,
    migration,
    plugins,
    profile,
    reports,
    settings,
    storage,
    tasks,
    translate,
)
from .api.routes.validation_tasks import router as validation_tasks
from .core.auth.security import get_password_hash
from ss_tools.shared.cot_logger import get_trace_id, seed_trace_id
from .core.database import AuthSessionLocal, init_db
from .core.encryption_key import ensure_encryption_key
from .core.logger import belief_scope, logger
from .core.utils.network import NetworkError
from .dependencies import get_async_job_runner, get_scheduler_service, get_task_manager
from .models.auth import Role, User


# #region lifespan [C:3] [TYPE Function]
# @ingroup Module
# @BRIEF Async context manager for FastAPI startup/shutdown lifecycle.
# @RELATION CALLS -> [init_db]
# @RELATION CALLS -> [AppDependencies]
# @POST On startup: admin exists, scheduler started. On shutdown: scheduler stopped.
# @RATIONALE Alembic migrations removed from lifespan — they now run exclusively
#   in docker/backend.entrypoint.sh (wait_for_db → alembic upgrade head).
#   Running migrations in both places added ~5s startup overhead and masked
#   partial failures. init_db() remains as a safety net for tables without
#   dedicated Alembic migrations (e.g., newly added models during development).
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    seed_trace_id()
    with belief_scope("startup_event"):
        logger.reason("Ensure encryption subsystem availability")
        ensure_encryption_key()
        logger.reason("Initialize persistent database tables")
        init_db()
        logger.reason("Bootstrap initial admin user (idempotent)")
        ensure_initial_admin_user()
        # Clean up stuck validation runs from previous backend lifetime
        # (runs left "running" in DB when the in-memory task queue was lost)
        try:
            from sqlalchemy.orm import Session as _Ses
            from src.core.database import SessionLocal as _Db
            from src.models.llm import ValidationRun as _VR
            from datetime import datetime, timezone

            _s: _Ses = _Db()
            _stuck = _s.query(_VR).filter(_VR.status == "running").all()
            for _r in _stuck:
                _r.status = "FAIL"
                _r.finished_at = datetime.now(timezone.utc)
                _r.summary = "Force-stopped: backend restarted while run was in progress"
                logger.reason(
                    "Force-stopped stuck run",
                    payload={"run_id": _r.id, "policy_id": _r.policy_id},
                )
            _s.commit()
            _s.close()
        except Exception as _e:
            logger.explore("Failed to clean up stuck validation runs", error=str(_e))

        # General reconciliation of stuck tasks from previous lifetime (idea from queue recovery patterns).
        # Tasks left in RUNNING when the in-memory asyncio tasks were lost (backend restart/crash).
        # This generalizes the ValidationRun-specific cleanup above.
        try:
            from sqlalchemy.orm import Session as _Ses
            from src.core.database import TasksSessionLocal as _TasksDb
            from src.models.task import TaskRecord as _TR
            from datetime import datetime as _dt, timezone as _tz

            _s: _Ses = _TasksDb()
            _stuck_tasks = _s.query(_TR).filter(_TR.status == "RUNNING").all()
            for _t in _stuck_tasks:
                _t.status = "FAILED"
                _t.finished_at = _dt.now(_tz.utc)
                _t.error = "Force-stopped: backend restarted while task was in progress"
                # result can carry hint
                if not _t.result:
                    _t.result = {"error": "interrupted_by_restart"}
                logger.reason(
                    "Force-stopped stuck task",
                    payload={"task_id": _t.id, "type": _t.type},
                )
            _s.commit()
            _s.close()
        except Exception as _e:
            logger.explore("Failed to clean up stuck general tasks", error=str(_e))

        logger.reason("Initializing AsyncJobRunner")
        get_async_job_runner()  # Initialize singleton with running event loop BEFORE scheduler starts
        logger.reason("Starting scheduler")
        scheduler = get_scheduler_service()
    scheduler.start()
    logger.reflect("Application startup complete")
    yield
    # Shutdown
    # Improved graceful shutdown (inspired by Celery worker drain + async best practices).
    # 1. Stop accepting new scheduled jobs.
    # 2. Attempt to let in-flight tasks finish or cancel them within timeout.
    # 3. Best-effort flush.
    scheduler.stop()
    try:
        from .dependencies import get_config_manager, get_task_manager
        tm = get_task_manager()
        # Access internal tracking (tasks are asyncio.Task objects)
        running = list(getattr(tm, '_async_tasks', {}).values())
        if running:
            cm = get_config_manager()
            graceful = getattr(cm, 'settings', None) if cm else None
            timeout = 30.0
            if graceful and hasattr(graceful, 'graceful_shutdown_timeout'):
                timeout = getattr(graceful, 'graceful_shutdown_timeout', 30.0)
            logger.reason(f"Draining {len(running)} running task(s) with timeout={timeout}s")
            done, pending = await asyncio.wait(running, timeout=timeout)
            for p in pending:
                p.cancel()
            if pending:
                await asyncio.gather(*pending, return_exceptions=True)
            logger.reason("Task drain complete")
    except Exception as _e:
        logger.explore("Graceful task drain during shutdown encountered error", error=str(_e))


# #endregion lifespan
# #region FastAPI_App [C:3] [TYPE Global] [SEMANTICS app, fastapi, instance, route-registry]
# @ingroup Module
# @BRIEF Canonical FastAPI application instance for route, middleware, and websocket registration.
# @RELATION DEPENDS_ON -> [ApiRoutesModule]
# @RELATION BINDS_TO -> [API_Routes]
app = FastAPI(
    title="Superset Tools API",
    description="API for managing Superset automation tools and plugins.",
    version="1.0.0",
    lifespan=lifespan,
)
# TraceContextMiddleware is a raw ASGI middleware (not BaseHTTPMiddleware).
# This is by design — BaseHTTPMiddleware (Starlette 0.50.0) uses
# anyio.create_task_group() internally, creating separate asyncio tasks for
# dispatch vs call_next. ContextVars set in dispatch() are NOT visible to
# outer middleware layers (like log_requests). Raw ASGI middleware runs in
# the root task context, making trace_id visible to ALL middleware layers.
# See trace.py @RATIONALE for details.
from .core.middleware.trace import TraceContextMiddleware  # noqa: E402

app.add_middleware(TraceContextMiddleware)


# #endregion FastAPI_App
# #region ensure_initial_admin_user [C:3] [TYPE Function]
# @ingroup Module
# @BRIEF Ensures initial admin user exists when bootstrap env flags are enabled.
# @RELATION DEPENDS_ON -> [AuthRepository]
def ensure_initial_admin_user() -> None:
    raw_flag = os.getenv("INITIAL_ADMIN_CREATE", "false").strip().lower()
    if raw_flag not in {"1", "true", "yes", "on"}:
        return
    username = os.getenv("INITIAL_ADMIN_USERNAME", "").strip()
    password = os.getenv("INITIAL_ADMIN_PASSWORD", "").strip()
    if not username or not password:
        logger.explore("INITIAL_ADMIN_CREATE enabled but credentials missing; skipping bootstrap")
        return
    # Warn about env-var password — visible via /proc to other processes
    logger.explore(
        "INITIAL_ADMIN_PASSWORD set via env var — visible to other processes",
        error="Security concern: password in environment variable",
    )
    db = AuthSessionLocal()
    try:
        admin_role = db.query(Role).filter(Role.name == "Admin").first()
        if not admin_role:
            admin_role = Role(name="Admin", description="System Administrator", is_admin=True)
            db.add(admin_role)
            db.commit()
            db.refresh(admin_role)
        existing_user = db.query(User).filter(User.username == username).first()
        if existing_user:
            logger.reflect(
                "Initial admin bootstrap skipped",
                payload={"username": username},
            )
            return
        new_user = User(
            username=username,
            email=None,
            password_hash=get_password_hash(password),
            auth_source="LOCAL",
            is_active=True,
        )
        new_user.roles.append(admin_role)
        db.add(new_user)
        db.commit()
        logger.reason(
            "Initial admin user created from environment bootstrap",
            payload={"username": username},
        )
    except Exception as exc:
        db.rollback()
        logger.explore("Failed to bootstrap initial admin user", error=str(exc))
        raise
    finally:
        db.close()


# #endregion ensure_initial_admin_user
# #region run_alembic_migrations [C:2] [TYPE Function]
# @ingroup Module
# @BRIEF Applies all pending Alembic migrations against DATABASE_URL.
#   DEPRECATED: Migrations now run exclusively in docker/backend.entrypoint.sh.
#   Kept for local development (manual invocation).
# @POST All Alembic migrations up to 'head' are applied.
# @SIDE_EFFECT Executes ALTER TABLE / CREATE TABLE via Alembic chain.
# @DEPRECATED Migrations moved to entrypoint.sh — this function is kept for
#   local development only. Do NOT call from lifespan or production code.
def run_alembic_migrations() -> None:
    from alembic.config import Config as AlembicConfig
    from alembic import command as alembic_command

    with belief_scope("run_alembic_migrations"):
        try:
            alembic_cfg = AlembicConfig("alembic.ini")
            alembic_cfg.set_main_option("script_location", "alembic")
            alembic_command.upgrade(alembic_cfg, "head")
            logger.reason("Alembic migrations applied up to head")
        except Exception as exc:
            logger.explore(
                "Alembic migration failed — check migration chain or alembic.ini",
                error=str(exc),
            )
            raise


# #endregion run_alembic_migrations


# #region app_middleware [TYPE Block]
# @ingroup Module
# @BRIEF Configure application-wide middleware (Session, CORS).
# @RATIONALE SessionMiddleware uses SESSION_SECRET_KEY independent of JWT SECRET_KEY (see [SEC:H-4]).
#            CORS allow_origins crashes early if ALLOWED_ORIGINS is unset — no "*" fallback (see [SEC:M-1]).
# @REJECTED Hardcoded allow_origins=["*"] rejected — open CORS allows any origin to access the API,
#           which is a Class 1 security violation in production.
# @REJECTED SessionMiddleware sharing JWT SECRET_KEY rejected in [SEC:H-4] — key reuse expands blast radius.

# Configure Session Middleware (required by Authlib for OAuth2 flow)
# Uses AUTH_SECRET_KEY directly — separate SESSION_SECRET_KEY eliminated.
from .core.auth.config import auth_config

app.add_middleware(SessionMiddleware, secret_key=auth_config.SECRET_KEY)

# Configure CORS
_allowed_origins_raw = os.getenv("ALLOWED_ORIGINS", "").strip()
if not _allowed_origins_raw:
    logger.explore(
        "ALLOWED_ORIGINS not set — CORS rejects all cross-origin requests",
        error="Missing ALLOWED_ORIGINS",
    )
    _allowed_origins = []
else:
    _allowed_origins = [o.strip() for o in _allowed_origins_raw.split(",") if o.strip()]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_allowed_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# HSTS — Strict-Transport-Security header (see [SEC:L-2])
# Only active when FORCE_HTTPS=true (enterprise deployments with proper certs).
# In dev or behind HTTP-only proxies, HSTS is disabled to avoid lockout.
class HSTSMiddleware(BaseHTTPMiddleware):
    """Add Strict-Transport-Security header when FORCE_HTTPS=true.

    Enterprise note: In production, set HSTS at the nginx/ingress level
    (add_header Strict-Transport-Security ...). This middleware is a
    fallback for environments where nginx is not configured to do so.
    """

    def __init__(self, app):
        super().__init__(app)
        self._enabled = os.getenv("FORCE_HTTPS", "").strip().lower() in {"1", "true", "yes"}

    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        if self._enabled:
            response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
        return response


app.add_middleware(HSTSMiddleware)


# #endregion app_middleware
# #region global_exception_handler [C:2] [TYPE Function]
# @ingroup Module
# @BRIEF Global exception handler — logs all unhandled 500 errors into the app logger.
# @PRE request is a FastAPI Request object.
# @POST Logs full traceback to superset_tools_app logger; returns 500.
# @RATIONALE FastAPI/Starlette writes unhandled exceptions to uvicorn.error logger,
#           which is not captured in docker logs. This handler ensures every 500
#           appears in our structured JSON log output.
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    client_host = request.client.host if request.client else "unknown"
    logger.explore(
        "Unhandled exception",
        payload={
            "method": request.method,
            "path": request.url.path,
            "query_params": dict(request.query_params),
            "client": client_host,
        },
        error=str(exc),
    )
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"detail": "Internal server error", "path": request.url.path},
    )


# #endregion global_exception_handler


# #region network_error_handler [C:1] [TYPE Function]
# @BRIEF Global exception handler for NetworkError.
# @PRE request is a FastAPI Request object.
# @POST Returns 503 HTTP Exception.
@app.exception_handler(NetworkError)
async def network_error_handler(request: Request, exc: NetworkError):
    with belief_scope("network_error_handler"):
        logger.explore("Network error", error=str(exc))
        return HTTPException(
            status_code=503,
            detail="Environment unavailable. Please check if the Superset instance is running.",
        )


# #endregion network_error_handler
# #region log_requests [C:3] [TYPE Function]
# @ingroup Module
# @BRIEF Middleware to log incoming HTTP requests and their response status.
# @RELATION DEPENDS_ON -> [LoggerModule]
# @PRE request is a FastAPI Request object.
# @POST Logs request and response details.
@app.middleware("http")
async def log_requests(request: Request, call_next):
    # Defensive trace_id seeding: TraceContextMiddleware (raw ASGI) seeds
    # trace_id in the root context, but BaseHTTPMiddleware may run in a
    # child task that doesn't inherit the ContextVar (Starlette 0.50.0 +
    # anyio.create_task_group). If empty, seed now.
    if not get_trace_id():
        seed_trace_id()
    with belief_scope("log_requests"):
        # Dynamic src derived from the request route — more informative than hardcoded "api.request_handler"
        _route_path = request.url.path.strip("/").replace("/", ".")
        _src = f"route.{request.method}.{_route_path}" if _route_path else f"route.{request.method}.root"
        is_polling = (
            (request.url.path.endswith("/api/tasks") and request.method == "GET")
            or request.url.path.endswith("/api/health/summary")
        )
        if not is_polling:
            logger.reason(
                "Handle API request",
                src=_src,
                payload={"method": request.method, "path": request.url.path}
            )
        try:
            response = await call_next(request)
            if not is_polling:
                logger.reflect(
                    "API request completed",
                    src=_src,
                    payload={"status": response.status_code, "path": request.url.path}
                )
            return response
        except NetworkError as e:
            logger.explore("Network error caught in middleware", error=str(e))
            raise HTTPException(
                status_code=503,
                detail="Environment unavailable. Please check if the Superset instance is running.",
            )


# #endregion log_requests
# #region API_Routes [C:3] [TYPE Block]
# @ingroup Module
# @BRIEF Register all FastAPI route groups exposed by the application entrypoint.
# @RELATION DEPENDS_ON -> [FastAPI_App]
# @RELATION DEPENDS_ON -> [Route_Group_Contracts]
# @RELATION DEPENDS_ON -> [Api.Auth]
# @RELATION DEPENDS_ON -> [AdminApi]
# @RELATION DEPENDS_ON -> [PluginsRouter]
# @RELATION DEPENDS_ON -> [TasksRouter]
# @RELATION DEPENDS_ON -> [SettingsRouter]
# @RELATION DEPENDS_ON -> [ReportsRouter]
# @RELATION DEPENDS_ON -> [LlmRoutes]
# @RELATION DEPENDS_ON -> [CleanReleaseV2Api]
# @RELATION DEPENDS_ON -> [MaintenanceRouter]
# Include API routes
app.include_router(auth.router)
app.include_router(admin.router)
app.include_router(admin_api_keys.router)
app.include_router(plugins.router, prefix="/api/plugins", tags=["Plugins"])
app.include_router(tasks.router, prefix="/api/tasks", tags=["Tasks"])
app.include_router(settings.router, prefix="/api/settings", tags=["Settings"])
app.include_router(environments.router, tags=["Environments"])
app.include_router(mappings.router, prefix="/api/mappings", tags=["Mappings"])
app.include_router(migration.router)
app.include_router(git.router, prefix="/api/git", tags=["Git"])
app.include_router(llm.router, prefix="/api/llm", tags=["LLM"])
app.include_router(storage.router, prefix="/api/storage", tags=["Storage"])
app.include_router(dashboards.router)
app.include_router(datasets.router)
app.include_router(reports.router)
app.include_router(assistant.router, prefix="/api/assistant", tags=["Assistant"])
app.include_router(agent_conversations.agent_router, tags=["Agent"])
app.include_router(agent_conversations.router, tags=["Assistant"])
app.include_router(agent_status.router)
app.include_router(agent_superset.router, tags=["Agent Superset"])
app.include_router(agent_superset_explore.router, tags=["Agent Superset"])
app.include_router(clean_release.router)
app.include_router(clean_release_v2.router)
app.include_router(profile.router)
app.include_router(dataset_review.router)
app.include_router(health.router)
app.include_router(encryption_health.router)
app.include_router(translate.router)
app.include_router(validation_tasks, prefix="/api/validation-tasks", tags=["Validation Tasks"])
app.include_router(maintenance.maintenance_router)


# #endregion API_Routes
# #region api.include_routers [C:1] [TYPE Action] [SEMANTICS routes, registration, api]
# @BRIEF Registers all API routers with the FastAPI application.
# @LAYER API
# #endregion api.include_routers
# #region _authenticate_websocket [TYPE Function]
# @BRIEF Authenticate a WebSocket connection via JWT or API key from query param `token`.
# @PRE  websocket is a live Starlette WebSocket before accept().
# @POST  Returns True if token is valid, logs reason; returns False if rejected.
# @RELATION DEPENDS_ON -> [Auth.Jwt]
# @RELATION DEPENDS_ON -> [APIKeyModel]
# @SIDE_EFFECT  Performs DB read to validate API key hash.
async def _authenticate_websocket(websocket: WebSocket, endpoint_name: str) -> bool:
    ws_token = websocket.query_params.get("token", "")
    if not ws_token:
        logger.explore(
            "WebSocket connection rejected",
            payload={"endpoint": endpoint_name},
            error="Missing token",
        )
        return False

    # Try JWT first, fallback to API key
    try:
        from .core.auth.api_key import hash_api_key
        from .core.auth.jwt import decode_token
        from .core.database import SessionLocal

        # Attempt JWT validation
        payload = decode_token(ws_token)
        user = payload.get("sub")
        if isinstance(user, str) and user:
            logger.reason(
                "WebSocket authenticated via JWT",
                payload={"endpoint": endpoint_name, "user": user},
            )
            return True
    except Exception:
        logger.explore("JWT validation failed in WebSocket auth", error="see traceback")
        pass

    # Fallback: try API key
    try:
        from .core.auth.api_key import hash_api_key
        from .core.database import SessionLocal
        from .models.api_key import APIKey

        key_hash = hash_api_key(ws_token)
        db = SessionLocal()
        try:
            api_key = db.query(APIKey).filter(APIKey.key_hash == key_hash).first()
            if api_key and api_key.active:
                logger.reason(
                    "WebSocket authenticated via API key",
                    payload={"endpoint": endpoint_name, "key_name": api_key.name},
                )
                return True
        finally:
            db.close()
    except Exception:
        logger.explore("API key validation failed in WebSocket auth", error="see traceback")
        pass

    logger.explore(
        "WebSocket connection rejected",
        payload={"endpoint": endpoint_name},
        error="Invalid token",
    )
    return False


# #endregion _authenticate_websocket


# #region websocket_endpoint [C:5] [TYPE Function]
# @ingroup Module
# @BRIEF Provides a WebSocket endpoint for real-time log streaming of a task with server-side filtering.
# @RELATION CALLS -> [TaskManagerPackage]
# @RELATION DEPENDS_ON -> [LoggerModule]
# @RELATION CALLS -> [_authenticate_websocket]
# @PRE task_id must be a valid task ID. WebSocket must be authenticated via `token` query param.
# @POST WebSocket connection is managed and logs are streamed until disconnect.
# @SIDE_EFFECT Subscribes to TaskManager log queue and broadcasts messages over network.
# @DATA_CONTRACT [task_id: str, source: str, level: str] -> [JSON log entry objects]
# @INVARIANT Every accepted WebSocket subscription is unsubscribed exactly once even when streaming fails or the client disconnects.
# @UX_STATE Connecting -> Streaming -> (Disconnected)
# @RATIONALE Uses asyncio.wait with two queues (log + status) to multiplex task logs and status
#           updates over a single WebSocket connection — eliminates the need for the client to
#           open a second socket or poll /api/tasks for status changes. Server-side source and
#           level filters reduce bandwidth by discarding irrelevant log entries before transmission.
# @REJECTED Separate WebSocket connections for logs vs status was rejected — doubles connection
#           overhead and complicates client-side synchronization. Client-side filtering was rejected
#           — sends unnecessary data over the network when the server can discard early. Polling
#           /api/tasks for status changes was rejected — introduces latency and load.
#
# @TEST_CONTRACT WebSocketLogStreamApi ->
# {
#   required_fields: {websocket: WebSocket, task_id: str},
#   optional_fields: {source: str, level: str},
#   invariants: [
#       "Accepts the WebSocket connection",
#       "Applies source and level filters correctly to streamed logs",
#       "Cleans up subscriptions on disconnect"
#   ]
# }
# @TEST_FIXTURE valid_ws_connection -> {"task_id": "test_1", "source": "plugin"}
# @TEST_EDGE task_not_found_ws -> closes connection or sends error
# @TEST_EDGE empty_task_logs -> waits for new logs
# @TEST_INVARIANT consistent_streaming -> verifies: [valid_ws_connection]
# @TEST_EDGE ws_auth_missing_token -> connection rejected with 4001
# @TEST_EDGE ws_auth_invalid_token -> connection rejected with 4001
@app.websocket("/ws/logs/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str, source: str = None, level: str = None):
    """
    WebSocket endpoint for real-time log streaming AND task status updates.
    Sends two message types:
      - Log entries: plain dicts with level/message/timestamp (backward compatible, no type field)
      - Status updates: dict with type="task_status" and nested task dict
    Query Parameters:
        source: Filter logs by source component (e.g., "plugin", "superset_api")
        level: Filter logs by minimum level (DEBUG, INFO, WARNING, ERROR)
        token: JWT or API key for authentication (required, see [SEC:C-3])
    """
    seed_trace_id()
    with belief_scope("websocket_endpoint", f"task_id={task_id}"):
        # ── WebSocket authentication (see [SEC:C-3]) ──
        if not await _authenticate_websocket(websocket, "ws/logs"):
            await websocket.close(code=4001, reason="Authentication required")
            return

        await websocket.accept()
        source_filter = source.lower() if source else None
        level_filter = level.upper() if level else None
        level_hierarchy = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
        min_level = level_hierarchy.get(level_filter, 0) if level_filter else 0
        logger.reason(
            "Accepted WebSocket log+status stream connection",
            payload={
                "task_id": task_id,
                "source_filter": source_filter,
                "level_filter": level_filter,
                "min_level": min_level,
            },
        )
        task_manager = get_task_manager()
        log_queue = await task_manager.subscribe_logs(task_id)
        status_queue = await task_manager.subscribe_status(task_id)
        logger.reason(
            "Subscribed WebSocket client to task log and status queues",
            payload={"task_id": task_id},
        )

        def matches_filters(log_entry) -> bool:
            """Check if log entry matches the filter criteria."""
            log_source = getattr(log_entry, "source", None)
            if source_filter and str(log_source or "").lower() != source_filter:
                return False
            if level_filter:
                log_level = level_hierarchy.get(str(log_entry.level).upper(), 0)
                if log_level < min_level:
                    return False
            return True

        async def send_status_update(task: Any) -> None:
            """Send a structured task status update over the WebSocket."""
            status_dict = {
                "id": task.id,
                "plugin_id": task.plugin_id,
                "status": task.status.value if hasattr(task.status, "value") else str(task.status),
                "started_at": task.started_at.isoformat() if task.started_at else None,
                "finished_at": task.finished_at.isoformat() if task.finished_at else None,
                "user_id": task.user_id,
                "result": task.result,
                "input_required": task.input_required,
                "input_request": task.input_request,
            }
            await websocket.send_json(
                {
                    "type": "task_status",
                    "task_id": task.id,
                    "task": status_dict,
                }
            )

        try:
            # ── Send initial task status ──
            task = task_manager.get_task(task_id)
            if task:
                await send_status_update(task)
                logger.reason(
                    "Sent initial task status",
                    payload={"task_id": task_id, "status": str(task.status)},
                )

            # ── Replay initial logs ──
            logger.reason(
                "Starting task log stream replay and live forwarding",
                payload={"task_id": task_id},
            )
            initial_logs = task_manager.get_task_logs(task_id)
            initial_sent = 0
            for log_entry in initial_logs:
                if matches_filters(log_entry):
                    log_dict = log_entry.model_dump()
                    log_dict["timestamp"] = log_dict["timestamp"].isoformat()
                    await websocket.send_json(log_dict)
                    initial_sent += 1
            logger.reflect(
                "Initial task log replay completed",
                payload={
                    "task_id": task_id,
                    "replayed_logs": initial_sent,
                    "total_available_logs": len(initial_logs),
                },
            )

            # ── Send synthetic AWAITING_INPUT prompt if needed ──
            task = task_manager.get_task(task_id)
            if task and task.status == "AWAITING_INPUT" and task.input_request:
                synthetic_log = {
                    "timestamp": task.logs[-1].timestamp.isoformat() if task.logs else "2024-01-01T00:00:00",
                    "level": "INFO",
                    "message": "Task paused for user input (Connection Re-established)",
                    "context": {"input_request": task.input_request},
                }
                await websocket.send_json(synthetic_log)
                logger.reason(
                    "Replayed awaiting-input prompt to restored WebSocket client",
                    payload={"task_id": task_id, "task_status": task.status},
                )

            # ── Main loop: listen on both log and status queues ──
            while True:
                log_task = asyncio.create_task(log_queue.get())
                status_task = asyncio.create_task(status_queue.get())
                done, _ = await asyncio.wait(
                    [log_task, status_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for coro in done:
                    result = coro.result()

                    # ── Status update ──
                    if isinstance(result, dict) and result.get("type") == "task_status":
                        await websocket.send_json(result)
                        task_status = result.get("task", {}).get("status", "")
                        if task_status in ("SUCCESS", "FAILED"):
                            logger.reason(
                                "Task reached terminal state via status broadcast; closing stream",
                                payload={"task_id": task_id, "status": task_status},
                            )
                            await asyncio.sleep(2)
                            raise StopIteration  # exit the while loop
                        continue

                    # ── Log entry ──
                    if not matches_filters(result):
                        continue
                    log_dict = result.model_dump()
                    log_dict["timestamp"] = log_dict["timestamp"].isoformat()
                    await websocket.send_json(log_dict)
                    logger.reflect(
                        "Forwarded task log entry to WebSocket client",
                        payload={
                            "task_id": task_id,
                            "level": log_dict.get("level"),
                        },
                    )
                    if "Task completed successfully" in result.message or "Task failed" in result.message:
                        logger.reason(
                            "Observed terminal task log entry; delaying to preserve client visibility",
                            payload={"task_id": task_id, "message": result.message},
                        )
                        await asyncio.sleep(2)
        except (WebSocketDisconnect, StopIteration) as _ws_exc:
            if isinstance(_ws_exc, StopIteration):
                # Task reached terminal state — close cleanly with code 1000
                try:
                    await websocket.close(code=1000, reason="Task completed")
                except Exception:
                    pass
            logger.reason(
                "WebSocket client disconnected or stream ended",
                payload={"task_id": task_id},
            )
        except Exception as exc:
            logger.explore(
                "WebSocket log+status streaming encountered an unexpected failure",
                payload={"task_id": task_id},
                error=str(exc),
            )
            raise
        finally:
            task_manager.unsubscribe_logs(task_id, log_queue)
            task_manager.unsubscribe_status(task_id, status_queue)
            logger.reflect(
                "Released WebSocket log and status queue subscriptions",
                payload={"task_id": task_id},
            )


# #endregion websocket_endpoint


# #region task_events_websocket [C:4] [TYPE Function]
# @ingroup Module
# @BRIEF WebSocket endpoint for global task events (status changes for ALL tasks).
# @RELATION CALLS -> [TaskManagerPackage]
# @PRE WebSocket must be authenticated via `token` query param.
# @POST WebSocket streams task status events until disconnect.
# @RATIONALE Push-based global task event stream replaces per-task status polling — a single
#           WebSocket delivers status changes for all tasks, enabling real-time dashboards and
#           notification UIs without repeated /api/tasks queries.
# @REJECTED Per-task WebSocket subscriptions was rejected — the task-manager already publishes
#           global events, so subscribing per-task duplicates the subscription overhead.
#           Polling /api/tasks was rejected — introduces latency, load on the DB, and
#           misses events between polls.
@app.websocket("/ws/task-events")
async def task_events_websocket(websocket: WebSocket):
    """
    WebSocket endpoint for global task events.
    Streams {type: "task_status", task_id: ..., task: {...}} for ALL task status changes.
    Query Parameters:
        token: JWT or API key for authentication (required)
    """
    seed_trace_id()
    with belief_scope("task_events_websocket"):
        if not await _authenticate_websocket(websocket, "ws/task-events"):
            await websocket.close(code=4001, reason="Authentication required")
            return

        await websocket.accept()
        logger.reason("Accepted global task events WebSocket connection")

        task_manager = get_task_manager()
        event_queue = await task_manager.subscribe_task_events()
        logger.reason("Subscribed to global task events")

        try:
            while True:
                event = await event_queue.get()
                await websocket.send_json(event)
                logger.reflect(
                    "Forwarded task event to global client",
                    payload={"task_id": event.get("task_id")},
                )
        except WebSocketDisconnect:
            logger.reason("Global task events client disconnected")
        except Exception as exc:
            logger.explore(
                "Global task events streaming failed",
                error=str(exc),
            )
            raise
        finally:
            task_manager.unsubscribe_task_events(event_queue)
            logger.reflect("Released global task events subscription")


# #endregion task_events_websocket


# #region maintenance_events_websocket [C:4] [TYPE Function]
# @ingroup Module
# @BRIEF WebSocket endpoint for maintenance events (created/ended/banner changes).
# @RELATION CALLS -> [TaskManagerPackage]
# @PRE WebSocket must be authenticated via `token` query param.
# @POST WebSocket streams maintenance events until disconnect.
# @RATIONALE Push-based maintenance event notification replaces polling /api/maintenance/events —
#           maintenance windows are time-sensitive (banner visibility, creation/ending); push
#           ensures sub-second delivery without DB load from polling.
# @REJECTED Polling /api/maintenance/events was rejected — maintenance state changes are
#           infrequent but time-critical; polling wastes resources during long idle periods
#           and still incurs latency.
@app.websocket("/ws/maintenance/events")
async def maintenance_events_websocket(websocket: WebSocket):
    """
    WebSocket endpoint for maintenance events.
    Streams {type: "maintenance.event_created", ...} / {type: "maintenance.event_ended"}.
    Query Parameters:
        token: JWT or API key for authentication (required)
    """
    seed_trace_id()
    with belief_scope("maintenance_events_websocket"):
        if not await _authenticate_websocket(websocket, "ws/maintenance/events"):
            await websocket.close(code=4001, reason="Authentication required")
            return

        await websocket.accept()
        logger.reason("Accepted maintenance events WebSocket connection")

        task_manager = get_task_manager()
        event_queue = await task_manager.subscribe_maintenance_events()
        logger.reason("Subscribed to maintenance events")

        try:
            while True:
                event = await event_queue.get()
                await websocket.send_json(event)
                logger.reflect(
                    "Forwarded maintenance event to client",
                    payload={"event_type": event.get("type"), "maintenance_id": event.get("maintenance_id")},
                )
        except WebSocketDisconnect:
            logger.reason("Maintenance events client disconnected")
        except Exception as exc:
            logger.explore(
                "Maintenance events streaming failed",
                error=str(exc),
            )
            raise
        finally:
            task_manager.unsubscribe_maintenance_events(event_queue)
            logger.reflect("Released maintenance events subscription")


# #endregion maintenance_events_websocket


# #region dataset_websocket_endpoint [C:4] [TYPE Function]
# @ingroup Module
# @BRIEF WebSocket endpoint for dataset.updated events — auto-refresh on task completion.
# @RELATION CALLS -> [TaskManagerPackage]
# @PRE env_id must reference a known environment.
# @POST WebSocket streams dataset.updated events until disconnect.
# @SIDE_EFFECT Subscribes to dataset event queue in task manager lifecycle.
@app.websocket("/ws/datasets/{env_id}")
async def dataset_websocket_endpoint(websocket: WebSocket, env_id: str):
    seed_trace_id()
    with belief_scope("dataset_websocket_endpoint", f"env_id={env_id}"):
        # ── WebSocket authentication (see [SEC:C-3]) ──
        if not await _authenticate_websocket(websocket, "ws/datasets"):
            await websocket.close(code=4001, reason="Authentication required")
            return

        await websocket.accept()
        logger.reason("Accepted dataset event WebSocket", payload={"env_id": env_id})
        task_manager = get_task_manager()
        queue = await task_manager.subscribe_dataset_events(env_id)
        try:
            while True:
                event = await queue.get()
                await websocket.send_json(event)
                logger.reflect("Forwarded dataset.updated event to client", payload={"env_id": env_id})
        except WebSocketDisconnect:
            logger.reason("WebSocket client disconnected from dataset events", payload={"env_id": env_id})
        except Exception as exc:
            logger.explore("WebSocket dataset streaming failed", payload={"env_id": env_id}, error=str(exc))
        finally:
            task_manager.unsubscribe_dataset_events(env_id, queue)
            logger.reflect("Released dataset event subscription", payload={"env_id": env_id})


# #endregion dataset_websocket_endpoint
# #region translate_run_websocket [C:3] [TYPE Function]
# @ingroup Module
# @BRIEF WebSocket endpoint for translation run progress — streams structured status updates.
# @PRE run_id must be a valid translation run ID. WebSocket authenticated via `token` query param.
# @POST Streams run status JSON every second until terminal state or disconnect.
# @SIDE_EFFECT Queries DB each tick for current run status via TranslationOrchestrator.
# @UX_STATE Streaming -> Terminal (completed/failed/cancelled) -> Close
# @UX_FEEDBACK Client receives {status, total_records, successful_records, failed_records, progressPct, ...}
@app.websocket("/ws/translate/run/{run_id}")
async def translate_run_websocket(websocket: WebSocket, run_id: str):
    seed_trace_id()
    if not await _authenticate_websocket(websocket, "ws/translate/run"):
        await websocket.close(code=4001, reason="Authentication required")
        return
    await websocket.accept()
    logger.reason("Accepted translate run WebSocket", payload={"run_id": run_id})
    try:
        while True:
            try:
                from .core.database import SessionLocal
                from .plugins.translate.orchestrator_aggregator import TranslationResultAggregator
                from .plugins.translate.events import TranslationEventLog

                db = SessionLocal()
                try:
                    event_log = TranslationEventLog(db)
                    aggregator = TranslationResultAggregator(db, event_log)
                    status = aggregator.get_run_status(run_id)
                    total = status.get("total_records", 0) or 0
                    done = (status.get("successful_records", 0) or 0) + (status.get("failed_records", 0) or 0) + (status.get("skipped_records", 0) or 0)
                    progress_pct = round((done / total) * 100) if total > 0 else 0
                    status["progressPct"] = progress_pct
                    await websocket.send_json(status)
                    if status.get("status") in ("COMPLETED", "FAILED", "CANCELLED"):
                        await asyncio.sleep(2)
                        break
                finally:
                    db.close()
            except Exception as tick_err:
                logger.explore("Translate run WS tick error", payload={"run_id": run_id}, error=str(tick_err))
                await websocket.send_json({"error": str(tick_err)})
                break
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        logger.reason("Translate run WS disconnected", payload={"run_id": run_id})
    except Exception as exc:
        logger.explore("Translate run WS error", payload={"run_id": run_id}, error=str(exc))
    logger.reflect("Translate run WS closed", payload={"run_id": run_id})


# #endregion translate_run_websocket
# #region StaticFiles [C:1] [TYPE Mount] [SEMANTICS static, frontend, spa]
# @BRIEF Mounts the frontend build directory to serve static assets.
frontend_path = project_root / "frontend" / "build"
if frontend_path.exists():
    app.mount("/_app", StaticFiles(directory=str(frontend_path / "_app")), name="static")

    # #region serve_spa [TYPE Function] [C:1]
    # @PURPOSE Serves the SPA frontend for any path not matched by API routes.
    # @PRE frontend_path exists.
    # @POST Returns the requested file or index.html.
    @app.get("/{file_path:path}", include_in_schema=False)
    async def serve_spa(file_path: str):
        with belief_scope("serve_spa"):
            # Only serve SPA for non-API paths
            # API routes are registered separately and should be matched by FastAPI first
            if file_path and (file_path.startswith("api/") or file_path.startswith("/api/") or file_path == "api"):
                # This should not happen if API routers are properly registered
                # Return 404 instead of serving HTML
                raise HTTPException(status_code=404, detail=f"API endpoint not found: {file_path}")
        full_path = frontend_path / file_path
        if file_path and full_path.is_file():
            return FileResponse(str(full_path))
        return FileResponse(str(frontend_path / "index.html"))

    # #endregion serve_spa
else:
    # #region read_root [TYPE Function] [C:1]
    # @PURPOSE A simple root endpoint to confirm that the API is running when frontend is missing.
    # @PRE None.
    # @POST Returns a JSON message indicating API status.
    @app.get("/")
    async def read_root():
        with belief_scope("read_root"):
            return {"message": "Superset Tools API is running (Frontend build not found)"}

    # #endregion read_root
# #endregion StaticFiles
# #endregion AppModule