1504 lines
65 KiB
Python
Executable File
1504 lines
65 KiB
Python
Executable File
# #region App.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 -> [Api.Init.ApiRoutesModule]
|
|
# @RELATION CALLS -> [Dependencies.AppDependencies.LiveExecutionComposition]
|
|
# @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]).
|
|
# @INVARIANT Startup initializes a fail-closed 044 LiveExecutionCompositionRoot; runtime providers
|
|
# must be registered server-side and never inferred from HTTP/run metadata.
|
|
# @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 uuid
|
|
|
|
# project_root is used for static files mounting
|
|
project_root = Path(__file__).resolve().parent.parent.parent
|
|
|
|
from fastapi import (
|
|
Depends,
|
|
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_lifecycle,
|
|
agent_runs,
|
|
agent_status,
|
|
agent_superset,
|
|
agent_superset_explore,
|
|
assistant,
|
|
clean_release,
|
|
clean_release_v2,
|
|
dashboard_testing,
|
|
dashboards,
|
|
datasets,
|
|
encryption_health,
|
|
environments,
|
|
git,
|
|
health,
|
|
llm,
|
|
maintenance,
|
|
mappings,
|
|
migration,
|
|
lineage,
|
|
load_testing,
|
|
plugins,
|
|
profile,
|
|
ready,
|
|
reports,
|
|
settings,
|
|
storage,
|
|
tasks,
|
|
tools,
|
|
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 build_cot_event, get_trace_id, seed_trace_id, set_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_config_manager,
|
|
get_current_user,
|
|
get_live_execution_composition_root,
|
|
get_scheduler_service,
|
|
get_task_manager,
|
|
)
|
|
from .models.auth import Role, User
|
|
|
|
|
|
# #region App.AppModule.LiveExecutionCompositionBootstrap [C:4] [TYPE Function] [SEMANTICS scenario,execution,composition,startup]
|
|
# @BRIEF Register trusted deployment live bindings before the scheduler can dispatch queued ScenarioRuns.
|
|
# @RELATION CALLS -> [ScenarioExecution.LiveCompositionRoot.Bootstrap]
|
|
# @POST Default dispatcher sees exact configured 037 providers or typed unavailable bindings; HTTP/run
|
|
# metadata never creates a live client.
|
|
# @INVARIANT Browser and Screenshot providers are separately registered server-side; their absence does
|
|
# not prevent startup and remains typed unavailable.
|
|
# @REJECTED Deferring Superset provider creation to a queued run was rejected — that would derive
|
|
# execution authority from a persisted/request identifier instead of deployment config.
|
|
def initialize_live_execution_composition() -> int:
|
|
from .services.dashboard_testing.execution.live_composition import (
|
|
bootstrap_live_execution_composition,
|
|
)
|
|
|
|
return bootstrap_live_execution_composition(
|
|
get_live_execution_composition_root(),
|
|
config_manager=get_config_manager(),
|
|
run_async=get_async_job_runner().run,
|
|
)
|
|
# #endregion App.AppModule.LiveExecutionCompositionBootstrap
|
|
|
|
|
|
# #region App.AppModule.Lifespan [C:3] [TYPE Function]
|
|
# @ingroup Module
|
|
# @BRIEF Async context manager for FastAPI startup/shutdown lifecycle.
|
|
# @RELATION CALLS -> [Core.Database.InitDb]
|
|
# @RELATION CALLS -> [Dependencies.AppDependencies]
|
|
# @RELATION CALLS -> [App.AppModule.LiveExecutionCompositionBootstrap]
|
|
# @POST On startup: admin exists, scheduler started. On shutdown: scheduler stopped.
|
|
# @INVARIANT The trusted live-composition bootstrap runs after AsyncJobRunner initialization and
|
|
# before scheduler startup; queued dispatch therefore cannot observe an unbootstrapped
|
|
# default root during normal application startup.
|
|
# @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("Synchronizing RBAC permission catalog at startup")
|
|
try:
|
|
from src.services.rbac_permission_catalog import (
|
|
discover_declared_permissions,
|
|
sync_permission_catalog,
|
|
)
|
|
from src.core.database import SessionLocal as _AuthDb
|
|
from src.dependencies import get_plugin_loader
|
|
from src.core.auth.repository import AuthRepository
|
|
|
|
_auth_db = _AuthDb()
|
|
_plugin_loader = get_plugin_loader()
|
|
_declared = discover_declared_permissions(plugin_loader=_plugin_loader)
|
|
_inserted = sync_permission_catalog(db=_auth_db, declared_permissions=_declared)
|
|
# Self-heal: guarantee the Admin role carries is_admin=True so role-based admin
|
|
# bypass works even on DBs created before the flag existed (see Fix #3).
|
|
AuthRepository(_auth_db).ensure_admin_role()
|
|
_auth_db.close()
|
|
if _inserted > 0:
|
|
logger.reason(
|
|
f"Synchronized {_inserted} new RBAC permissions at startup",
|
|
payload={"inserted": _inserted},
|
|
)
|
|
except Exception as _e:
|
|
logger.explore("Failed to sync RBAC permission catalog at startup", error=str(_e))
|
|
|
|
logger.reason("Initializing AsyncJobRunner")
|
|
get_async_job_runner() # Initialize singleton with running event loop BEFORE scheduler starts
|
|
logger.reason("Bootstrapping fail-closed ScenarioRun live composition root")
|
|
live_bindings = initialize_live_execution_composition()
|
|
logger.reason("ScenarioRun live composition bootstrap complete", payload={"bindings": live_bindings})
|
|
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 App.AppModule.Lifespan
|
|
# #region App.AppModule.FastAPIApp [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 -> [Api.Init.ApiRoutesModule]
|
|
# @RELATION BINDS_TO -> [App.AppModule.APIRoutes]
|
|
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 App.AppModule.FastAPIApp
|
|
# #region App.AppModule.EnsureInitialAdminUser [C:3] [TYPE Function]
|
|
# @ingroup Module
|
|
# @BRIEF Ensures initial admin user exists when bootstrap env flags are enabled.
|
|
# @RELATION DEPENDS_ON -> [Core.Repository.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 App.AppModule.EnsureInitialAdminUser
|
|
# #region App.AppModule.RunAlembicMigrations [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 App.AppModule.RunAlembicMigrations
|
|
|
|
|
|
# #region App.AppModule.AppMiddleware [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.AppModule.AppMiddleware
|
|
# #region App.AppModule.GlobalExceptionHandler [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 App.AppModule.GlobalExceptionHandler
|
|
|
|
|
|
# #region App.AppModule.NetworkErrorHandler [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 JSONResponse(
|
|
status_code=503,
|
|
content={
|
|
"detail": "Environment unavailable. Please check if the Superset instance is running.",
|
|
"path": request.url.path,
|
|
},
|
|
)
|
|
|
|
|
|
# #endregion App.AppModule.NetworkErrorHandler
|
|
|
|
# High-frequency polling endpoints whose per-request REASON/REFLECT framing would
|
|
# spam the structured log (they remain visible in the uvicorn access log):
|
|
# - /api/tasks* task progress polling (every 1.5s during operations)
|
|
# - /api/health/summary health monitoring polling
|
|
# - /api/agent/llm-config agent container LLM config polling
|
|
# - /api/auth/session/activity session activity heartbeats
|
|
# - /api/settings/consolidated settings polling
|
|
_POLLING_EXACT_PATHS = frozenset({
|
|
"/api/health/summary",
|
|
"/api/agent/llm-config",
|
|
"/api/auth/session/activity",
|
|
"/api/settings/consolidated",
|
|
})
|
|
|
|
|
|
def _is_suppressed_request(request: Request) -> bool:
|
|
"""Return True for high-frequency polling requests (framing suppressed)."""
|
|
path = request.url.path
|
|
return path in _POLLING_EXACT_PATHS or (
|
|
request.method == "GET" and path.startswith("/api/tasks")
|
|
)
|
|
|
|
|
|
# #region App.AppModule.LogRequests [C:3] [TYPE Function]
|
|
# @ingroup Module
|
|
# @BRIEF Middleware to log incoming HTTP requests and their response status.
|
|
# @RELATION DEPENDS_ON -> [Core.Logger.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 = _is_suppressed_request(request)
|
|
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 App.AppModule.LogRequests
|
|
# #region App.AppModule.APIRoutes [C:3] [TYPE Block]
|
|
# @ingroup Module
|
|
# @BRIEF Register all FastAPI route groups exposed by the application entrypoint.
|
|
# @RELATION DEPENDS_ON -> [App.AppModule.FastAPIApp]
|
|
# @RELATION DEPENDS_ON -> [Api.Init.RouteGroupContracts]
|
|
# @RELATION DEPENDS_ON -> [Api.Auth]
|
|
# @RELATION DEPENDS_ON -> [Api.Admin.AdminApi]
|
|
# @RELATION DEPENDS_ON -> [Api.Plugins.PluginsRouter]
|
|
# @RELATION DEPENDS_ON -> [Api.Tasks.TasksRouter]
|
|
# @RELATION DEPENDS_ON -> [Api.Settings.SettingsRouter]
|
|
# @RELATION DEPENDS_ON -> [Api.Reports.ReportsRouter]
|
|
# @RELATION DEPENDS_ON -> [Api.Llm.LlmRoutes]
|
|
# @RELATION DEPENDS_ON -> [Api.CleanReleaseV2.CleanReleaseV2Api]
|
|
# @RELATION DEPENDS_ON -> [Api.Router.MaintenanceRouter]
|
|
# @RELATION DEPENDS_ON -> [Api.DashboardTesting]
|
|
# 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(lineage.router)
|
|
app.include_router(load_testing.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(dashboard_testing.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_lifecycle.router)
|
|
app.include_router(agent_runs.router)
|
|
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(health.router)
|
|
app.include_router(encryption_health.router)
|
|
app.include_router(ready.router)
|
|
app.include_router(translate.router)
|
|
app.include_router(validation_tasks, prefix="/api/validation-tasks", tags=["Validation Tasks"])
|
|
|
|
|
|
# #region App.AppModule.GetRecentAppLogs [C:3] [TYPE Function] [SEMANTICS api,logs,app,recent]
|
|
# @BRIEF REST snapshot of the process-wide app/cot log ring buffer (raw JSON lines).
|
|
@app.get("/api/logs/recent")
|
|
async def get_recent_app_logs(
|
|
limit: int = 500,
|
|
level: str | None = None,
|
|
task_id: str | None = None,
|
|
after_seq: int = 0,
|
|
_user=Depends(get_current_user),
|
|
):
|
|
"""Return recent application log lines for the Reports live console.
|
|
|
|
Each item: {seq, level, raw, timestamp, logger}.
|
|
`raw` is CotJsonFormatter single-line JSON (JSONL-compatible).
|
|
Optional task_id: comma-separated filter — lines whose raw text contains the id.
|
|
"""
|
|
from src.core.logger import get_app_log_handler
|
|
|
|
limit = max(1, min(int(limit or 500), 3000))
|
|
handler = get_app_log_handler()
|
|
level_hierarchy = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
|
|
min_level = level_hierarchy.get((level or "").upper(), 0) if level else 0
|
|
task_ids = {t.strip() for t in (task_id or "").split(",") if t.strip()}
|
|
|
|
if after_seq > 0:
|
|
lines, current = handler.get_since(after_seq, limit=limit)
|
|
else:
|
|
lines = handler.get_recent_logs(limit=limit)
|
|
current = handler.current_seq()
|
|
|
|
items = []
|
|
for line in lines:
|
|
lv = level_hierarchy.get(str(line.level).upper(), 0)
|
|
if lv < min_level:
|
|
continue
|
|
raw = line.raw or ""
|
|
if task_ids and not any(tid in raw for tid in task_ids):
|
|
continue
|
|
items.append(
|
|
{
|
|
"seq": line.seq,
|
|
"level": line.level,
|
|
"raw": raw,
|
|
"timestamp": line.timestamp.isoformat() if line.timestamp else None,
|
|
"logger": line.logger_name,
|
|
}
|
|
)
|
|
return {"items": items, "current_seq": current, "capacity": handler.capacity}
|
|
|
|
|
|
# #endregion App.AppModule.GetRecentAppLogs
|
|
app.include_router(tools.router, tags=["Tools"])
|
|
app.include_router(maintenance.maintenance_router)
|
|
|
|
|
|
# #endregion App.AppModule.APIRoutes
|
|
# #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 App.AppModule.AuthenticateWebsocket [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 -> [Models.ApiKey.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 as exc:
|
|
logger.explore(
|
|
"JWT validation failed in WebSocket auth",
|
|
payload={"endpoint": endpoint_name, "reason": f"{type(exc).__name__}: {exc}"},
|
|
error="Invalid token",
|
|
)
|
|
|
|
# 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 as exc:
|
|
logger.explore(
|
|
"API key validation failed in WebSocket auth",
|
|
payload={"endpoint": endpoint_name, "reason": f"{type(exc).__name__}: {exc}"},
|
|
error="Invalid token",
|
|
)
|
|
|
|
logger.explore(
|
|
"WebSocket connection rejected",
|
|
payload={"endpoint": endpoint_name},
|
|
error="Invalid token",
|
|
)
|
|
return False
|
|
|
|
|
|
# #endregion App.AppModule.AuthenticateWebsocket
|
|
|
|
|
|
# #region App.AppModule.RejectWebsocketAuth [C:2] [TYPE Function] [SEMANTICS websocket,auth,close]
|
|
# @BRIEF Accept then close a WebSocket with the auth-rejection close code so clients can
|
|
# detect the failure via event.code instead of an opaque HTTP 403 handshake rejection.
|
|
# @POST Connection is accepted and immediately closed with code 4001.
|
|
async def _reject_websocket_auth(websocket: WebSocket) -> None:
|
|
try:
|
|
await websocket.accept()
|
|
await websocket.close(code=4001, reason="Authentication required")
|
|
except WebSocketDisconnect:
|
|
pass # Client already gone — nothing to send.
|
|
|
|
|
|
# #endregion App.AppModule.RejectWebsocketAuth
|
|
|
|
|
|
# #region App.AppModule.RejectWebsocketPermission [C:2] [TYPE Function] [SEMANTICS websocket,auth,rbac,close]
|
|
# @BRIEF Accept then close a WebSocket with the permission-rejection close code so clients can
|
|
# detect the failure via event.code instead of an opaque HTTP 403 handshake rejection.
|
|
# @POST Connection is accepted and immediately closed with code 4003.
|
|
async def _reject_websocket_permission(websocket: WebSocket) -> None:
|
|
try:
|
|
await websocket.accept()
|
|
await websocket.close(code=4003, reason="Insufficient permissions")
|
|
except WebSocketDisconnect:
|
|
pass # Client already gone — nothing to send.
|
|
|
|
|
|
# #endregion App.AppModule.RejectWebsocketPermission
|
|
|
|
|
|
|
|
# #region App.AppModule.AuthorizeWebsocket [C:3] [TYPE Function] [SEMANTICS websocket,auth,rbac]
|
|
# @ingroup Module
|
|
# @BRIEF Extract the authenticated principal from a WebSocket token and check RBAC permissions.
|
|
# @PRE _authenticate_websocket has already validated the token.
|
|
# @POST Returns True if the authenticated user has the required permission; False otherwise.
|
|
# @RELATION DEPENDS_ON -> [Auth.Jwt.DecodeToken]
|
|
# @RELATION DEPENDS_ON -> [Core.Repository.AuthRepository]
|
|
def _authorize_websocket(websocket: WebSocket, resource: str, action: str) -> bool:
|
|
"""Check if the WebSocket-authenticated user has the required RBAC permission."""
|
|
ws_token = websocket.query_params.get("token", "")
|
|
|
|
try:
|
|
from .core.auth.jwt import decode_token
|
|
from .core.database import SessionLocal
|
|
from .models.auth import User
|
|
|
|
payload = decode_token(ws_token)
|
|
username = payload.get("sub")
|
|
if not isinstance(username, str) or not username:
|
|
logger.explore(
|
|
"WebSocket authorization — token missing 'sub' claim",
|
|
payload={"resource": resource, "action": action},
|
|
error="JWT payload has no valid subject",
|
|
)
|
|
return False
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
user = db.query(User).filter(User.username == username).first()
|
|
if not user:
|
|
logger.explore(
|
|
"WebSocket authorization — user not found in database",
|
|
payload={"resource": resource, "action": action, "username": username},
|
|
error="Authenticated user missing from local DB",
|
|
)
|
|
return False
|
|
if not getattr(user, "is_active", True):
|
|
logger.explore(
|
|
"WebSocket authorization — user is inactive",
|
|
payload={"resource": resource, "action": action, "username": username},
|
|
error="User account is disabled",
|
|
)
|
|
return False
|
|
|
|
# is_admin is the sole administrative authority. The database migration
|
|
# backfills the legacy Admin role before this authorization path runs.
|
|
for role in user.roles:
|
|
if getattr(role, "is_admin", False):
|
|
return True
|
|
for perm in role.permissions:
|
|
if perm.resource == resource and perm.action == action:
|
|
return True
|
|
|
|
logger.explore(
|
|
"WebSocket authorization denied — no matching permission or admin role",
|
|
payload={"resource": resource, "action": action, "user": username,
|
|
"roles": [r.name for r in user.roles],
|
|
"is_active": getattr(user, "is_active", True)},
|
|
error="No matching RBAC permission",
|
|
)
|
|
return False
|
|
finally:
|
|
db.close()
|
|
except Exception:
|
|
logger.explore(
|
|
"WebSocket authorization failed",
|
|
payload={"resource": resource, "action": action},
|
|
error="Token decode or DB lookup failed",
|
|
)
|
|
return False
|
|
|
|
|
|
# #endregion App.AppModule.AuthorizeWebsocket
|
|
|
|
|
|
# #region App.AppModule.SetWebsocketTraceId [C:2] [TYPE Function] [SEMANTICS websocket,trace,context]
|
|
# @ingroup Module
|
|
# @BRIEF Apply a valid UUID4 x-trace-id query parameter to the current WebSocket context.
|
|
# @PRE websocket is a live WebSocket whose query params may include x-trace-id.
|
|
# @POST Valid UUID4 query value is propagated; invalid or missing values receive a new trace ID.
|
|
# @SIDE_EFFECT Sets the shared trace ContextVar for this WebSocket task.
|
|
# @RATIONALE Browser WebSocket clients cannot set arbitrary request headers, so trace
|
|
# propagation uses a query parameter alongside the authentication token.
|
|
# @REJECTED Requiring a custom WebSocket header was rejected — browser WebSocket APIs do not
|
|
# permit application-defined headers during the opening handshake.
|
|
def _set_websocket_trace_id(websocket: WebSocket) -> str:
|
|
incoming = websocket.query_params.get("x-trace-id", "")
|
|
if incoming:
|
|
try:
|
|
parsed = uuid.UUID(hex=incoming)
|
|
if parsed.version == 4:
|
|
set_trace_id(incoming)
|
|
return incoming
|
|
except (ValueError, AttributeError):
|
|
pass
|
|
return seed_trace_id()
|
|
|
|
|
|
# #endregion App.AppModule.SetWebsocketTraceId
|
|
|
|
|
|
# #region App.AppModule.WebsocketEndpoint [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 -> [Core.Init.TaskManagerPackage]
|
|
# @RELATION DEPENDS_ON -> [Core.Logger.LoggerModule]
|
|
# @RELATION CALLS -> [App.AppModule.AuthenticateWebsocket]
|
|
# @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])
|
|
"""
|
|
_set_websocket_trace_id(websocket)
|
|
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 _reject_websocket_auth(websocket)
|
|
return
|
|
|
|
if not _authorize_websocket(websocket, "tasks", "READ"):
|
|
await _reject_websocket_permission(websocket)
|
|
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, "src", None) or 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 = build_cot_event(
|
|
src="task.lifecycle.websocket_reconnect",
|
|
marker="REFLECT",
|
|
intent="Task paused for user input (connection re-established)",
|
|
payload={"input_request": task.input_request},
|
|
)
|
|
synthetic_log = {"task_id": task_id, **synthetic_log}
|
|
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, pending = await asyncio.wait(
|
|
[log_task, status_task],
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
# @INVARIANT Cancel the not-yet-completed waiter BEFORE the next iteration.
|
|
# Without this, each loop leaks a pending queue.get() coroutine that
|
|
# consumes the NEXT item (e.g. the terminal task_status event), so the
|
|
# completion broadcast is silently swallowed and the client never sees
|
|
# the structured result — the result panel stays hidden after a task ends.
|
|
for pending_task in pending:
|
|
pending_task.cancel()
|
|
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"),
|
|
},
|
|
)
|
|
result_intent = getattr(result, "intent", None) or getattr(result, "message", "")
|
|
if "Task completed successfully" in result_intent or "Task failed" in result_intent:
|
|
logger.reason(
|
|
"Observed terminal task log entry; delaying to preserve client visibility",
|
|
payload={"task_id": task_id, "intent": result_intent},
|
|
)
|
|
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 App.AppModule.WebsocketEndpoint
|
|
|
|
|
|
# #region App.AppModule.TaskEventsWebsocket [C:4] [TYPE Function]
|
|
# @ingroup Module
|
|
# @BRIEF WebSocket endpoint for global task events (status changes for ALL tasks).
|
|
# @RELATION CALLS -> [Core.Init.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)
|
|
"""
|
|
_set_websocket_trace_id(websocket)
|
|
with belief_scope("task_events_websocket"):
|
|
if not await _authenticate_websocket(websocket, "ws/task-events"):
|
|
await _reject_websocket_auth(websocket)
|
|
return
|
|
|
|
if not _authorize_websocket(websocket, "tasks", "READ"):
|
|
await _reject_websocket_permission(websocket)
|
|
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 App.AppModule.TaskEventsWebsocket
|
|
|
|
|
|
# #region App.AppModule.AppLogsWebsocket [C:4] [TYPE Function] [SEMANTICS websocket,logs,app,realtime,jsonl]
|
|
# @Namespace Module
|
|
# @BRIEF Live tail of process-wide app/cot log ring buffer (raw JSON lines).
|
|
# @PRE WebSocket authenticated via token query param.
|
|
# @POST Streams {seq, level, raw, timestamp} for new log lines until disconnect.
|
|
# @RATIONALE Reports Logs tab needs full server stream, not only task-scoped logs.
|
|
# @REJECTED Binding console solely to task_id — operators need unfiltered live server output.
|
|
@app.websocket("/ws/app-logs")
|
|
async def app_logs_websocket(
|
|
websocket: WebSocket,
|
|
level: str | None = None,
|
|
task_id: str | None = None,
|
|
):
|
|
"""
|
|
Stream application CoT/JSON log lines from the in-process ring buffer.
|
|
Query params:
|
|
token: JWT (required)
|
|
level: min level (DEBUG|INFO|WARNING|ERROR)
|
|
task_id: optional comma-separated task ids — only lines whose raw JSON contains matching task_id
|
|
"""
|
|
import asyncio
|
|
|
|
from src.core.logger import get_app_log_handler
|
|
|
|
_set_websocket_trace_id(websocket)
|
|
with belief_scope("app_logs_websocket"):
|
|
if not await _authenticate_websocket(websocket, "ws/app-logs"):
|
|
await _reject_websocket_auth(websocket)
|
|
return
|
|
|
|
if not _authorize_websocket(websocket, "admin:settings", "READ"):
|
|
await _reject_websocket_permission(websocket)
|
|
return
|
|
|
|
await websocket.accept()
|
|
handler = get_app_log_handler()
|
|
level_hierarchy = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
|
|
min_level = level_hierarchy.get((level or "").upper(), 0) if level else 0
|
|
task_ids = {t.strip() for t in (task_id or "").split(",") if t.strip()}
|
|
|
|
def _match(line) -> bool:
|
|
lv = level_hierarchy.get(str(line.level).upper(), 0)
|
|
if lv < min_level:
|
|
return False
|
|
if not task_ids:
|
|
return True
|
|
raw = line.raw or ""
|
|
# Fast path: substring match for task uuid in JSON
|
|
return any(tid in raw for tid in task_ids)
|
|
|
|
# Snapshot recent, then poll for new seq
|
|
recent = handler.get_recent_logs(limit=500)
|
|
last_seq = handler.current_seq()
|
|
for line in recent:
|
|
if _match(line):
|
|
await websocket.send_json({
|
|
"seq": line.seq,
|
|
"level": line.level,
|
|
"raw": line.raw,
|
|
"timestamp": line.timestamp.isoformat() if line.timestamp else None,
|
|
"logger": line.logger_name,
|
|
})
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(0.25)
|
|
newer, last_seq = handler.get_since(last_seq, limit=500)
|
|
for line in newer:
|
|
if _match(line):
|
|
await websocket.send_json({
|
|
"seq": line.seq,
|
|
"level": line.level,
|
|
"raw": line.raw,
|
|
"timestamp": line.timestamp.isoformat() if line.timestamp else None,
|
|
"logger": line.logger_name,
|
|
})
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception as exc:
|
|
logger.explore("App log stream failed", error=str(exc))
|
|
raise
|
|
|
|
|
|
# #endregion App.AppModule.AppLogsWebsocket
|
|
|
|
|
|
|
|
# #region App.AppModule.MaintenanceEventsWebsocket [C:4] [TYPE Function]
|
|
# @ingroup Module
|
|
# @BRIEF WebSocket endpoint for maintenance events (created/ended/banner changes).
|
|
# @RELATION CALLS -> [Core.Init.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)
|
|
"""
|
|
_set_websocket_trace_id(websocket)
|
|
with belief_scope("maintenance_events_websocket"):
|
|
if not await _authenticate_websocket(websocket, "ws/maintenance/events"):
|
|
await _reject_websocket_auth(websocket)
|
|
return
|
|
|
|
if not _authorize_websocket(websocket, "maintenance", "READ"):
|
|
await _reject_websocket_permission(websocket)
|
|
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 App.AppModule.MaintenanceEventsWebsocket
|
|
|
|
|
|
# #region App.AppModule.DatasetWebsocketEndpoint [C:4] [TYPE Function]
|
|
# @ingroup Module
|
|
# @BRIEF WebSocket endpoint for dataset.updated events — auto-refresh on task completion.
|
|
# @RELATION CALLS -> [Core.Init.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):
|
|
_set_websocket_trace_id(websocket)
|
|
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 _reject_websocket_auth(websocket)
|
|
return
|
|
|
|
if not _authorize_websocket(websocket, "plugin:migration", "READ"):
|
|
await _reject_websocket_permission(websocket)
|
|
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 App.AppModule.DatasetWebsocketEndpoint
|
|
# #region App.AppModule.TranslateRunWebsocket [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, ...}
|
|
def _ws_is_client_gone(exc: BaseException) -> bool:
|
|
"""True when the peer already closed the socket (do not send again)."""
|
|
if isinstance(exc, WebSocketDisconnect):
|
|
return True
|
|
if isinstance(exc, RuntimeError):
|
|
msg = str(exc).lower()
|
|
return "close message has been sent" in msg or "not connected" in msg
|
|
return False
|
|
|
|
|
|
@app.websocket("/ws/translate/run/{run_id}")
|
|
async def translate_run_websocket(websocket: WebSocket, run_id: str):
|
|
_set_websocket_trace_id(websocket)
|
|
if not await _authenticate_websocket(websocket, "ws/translate/run"):
|
|
await _reject_websocket_auth(websocket)
|
|
return
|
|
|
|
if not _authorize_websocket(websocket, "translate.run", "VIEW"):
|
|
await _reject_websocket_permission(websocket)
|
|
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 WebSocketDisconnect:
|
|
# Client gone mid-tick — not a status/DB failure; outer handler logs disconnect.
|
|
raise
|
|
except Exception as tick_err:
|
|
err_type = type(tick_err).__name__
|
|
err_msg = str(tick_err) or repr(tick_err)
|
|
if _ws_is_client_gone(tick_err):
|
|
logger.reason(
|
|
"Translate run WS client gone during tick",
|
|
payload={"run_id": run_id, "error_type": err_type},
|
|
error=err_msg,
|
|
)
|
|
break
|
|
logger.explore(
|
|
"Translate run WS tick error",
|
|
payload={"run_id": run_id, "error_type": err_type},
|
|
error=err_msg,
|
|
)
|
|
try:
|
|
await websocket.send_json({"error": err_msg, "error_type": err_type})
|
|
except Exception as send_err:
|
|
# Peer closed between tick failure and error frame — avoid second failure noise.
|
|
logger.reason(
|
|
"Translate run WS skip error frame (socket closed)",
|
|
payload={
|
|
"run_id": run_id,
|
|
"error_type": type(send_err).__name__,
|
|
},
|
|
error=str(send_err) or repr(send_err),
|
|
)
|
|
break
|
|
await asyncio.sleep(1)
|
|
except WebSocketDisconnect:
|
|
logger.reason("Translate run WS disconnected", payload={"run_id": run_id})
|
|
except Exception as exc:
|
|
err_type = type(exc).__name__
|
|
err_msg = str(exc) or repr(exc)
|
|
if _ws_is_client_gone(exc):
|
|
logger.reason(
|
|
"Translate run WS client gone",
|
|
payload={"run_id": run_id, "error_type": err_type},
|
|
error=err_msg,
|
|
)
|
|
else:
|
|
logger.explore(
|
|
"Translate run WS error",
|
|
payload={"run_id": run_id, "error_type": err_type},
|
|
error=err_msg,
|
|
)
|
|
logger.reflect("Translate run WS closed", payload={"run_id": run_id})
|
|
|
|
|
|
# #endregion App.AppModule.TranslateRunWebsocket
|
|
# #region App.AppModule.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 App.AppModule.ServeSpa [TYPE Function] [C:1]
|
|
# @BRIEF 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 App.AppModule.ServeSpa
|
|
else:
|
|
# #region App.AppModule.ReadRoot [TYPE Function] [C:1]
|
|
# @BRIEF 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 App.AppModule.ReadRoot
|
|
# #endregion App.AppModule.StaticFiles
|
|
# #endregion App.AppModule
|