feat: security headers, frontend auth context, ML registry, and orchestration pipeline

backend/core/middleware.py

@@ -1,7 +1,45 @@
 from fastapi import Request, Response
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, RedirectResponse
+from starlette.middleware.base import BaseHTTPMiddleware
 from slowapi import Limiter
 from slowapi.errors import RateLimitExceeded
+from backend.core.config import settings
+
+class HTTPSRedirectMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        if settings.is_production:
+            # Check if request came in as HTTP (Render.com forwards as HTTPS but sets X-Forwarded-Proto)
+            proto = request.headers.get("X-Forwarded-Proto", "https")
+            if proto == "http":
+                https_url = str(request.url).replace("http://", "https://", 1)
+                return RedirectResponse(https_url, status_code=301)
+        return await call_next(request)
+
+class SecurityHeadersMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        response = await call_next(request)
+        
+        # Always add these headers:
+        response.headers["X-Content-Type-Options"] = "nosniff"
+        response.headers["X-Frame-Options"] = "DENY"
+        response.headers["X-XSS-Protection"] = "1; mode=block"
+        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
+        response.headers["Permissions-Policy"] = (
+            "camera=(), microphone=(self), geolocation=(), "
+            "payment=(), usb=(), magnetometer=()"
+        ) # microphone=(self) required for Whisper voice input
+        
+        # Production only:
+        if settings.is_production:
+            response.headers["Strict-Transport-Security"] = (
+                "max-age=63072000; includeSubDomains; preload"
+            )
+        
+        # Remove headers that leak server info:
+        response.headers.pop("Server", None)
+        response.headers.pop("X-Powered-By", None)
+        
+        return response
 
 def get_rate_limit_key(request: Request) -> str:
     user_id = getattr(request.state, "user_id", None)

backend/ml/registry.py

@@ -0,0 +1,291 @@
+import asyncio
+import time
+import logging
+from dataclasses import dataclass
+from datetime import datetime, UTC
+from typing import Literal, Any
+import torch
+
+from backend.core.config import settings
+from backend.core.logging_config import ml_logger
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class ModelProfile:
+    name: str
+    hf_model_id: str
+    local_cache_subdir: str
+    device_preference: Literal["cuda", "cpu", "auto"]
+    vram_mb: int
+    ram_mb: int
+    load_priority: int
+    is_required: bool
+
+MODEL_PROFILES = {
+    "dino_anomaly": ModelProfile(
+        name="dino_anomaly", hf_model_id="facebook/dinov2-small", local_cache_subdir="dino_anomaly",
+        device_preference="cuda", vram_mb=400, ram_mb=50, load_priority=1, is_required=True
+    ),
+    "biobert_ner": ModelProfile(
+        name="biobert_ner", hf_model_id="dmis-lab/biobert-base-cased-v1.2", local_cache_subdir="biobert_ner",
+        device_preference="cuda", vram_mb=450, ram_mb=50, load_priority=2, is_required=True
+    ),
+    "biomedvlp": ModelProfile(
+        name="biomedvlp", hf_model_id="microsoft/BiomedVLP-CXR-BERT-specialized", local_cache_subdir="biomedvlp",
+        device_preference="auto", vram_mb=900, ram_mb=100, load_priority=3, is_required=False
+    ),
+    "whisper_tiny": ModelProfile(
+        name="whisper_tiny", hf_model_id="openai/whisper-tiny", local_cache_subdir="whisper",
+        device_preference="cpu", vram_mb=0, ram_mb=300, load_priority=4, is_required=False
+    ),
+    "biogpt_base": ModelProfile(
+        name="biogpt_base", hf_model_id="microsoft/biogpt", local_cache_subdir="biogpt",
+        device_preference="cpu", vram_mb=0, ram_mb=700, load_priority=5, is_required=False
+    ),
+    "minilm": ModelProfile(
+        name="minilm", hf_model_id="sentence-transformers/all-MiniLM-L6-v2", local_cache_subdir="minilm",
+        device_preference="cpu", vram_mb=0, ram_mb=100, load_priority=1, is_required=True
+    ),
+}
+
+@dataclass
+class ModelState:
+    profile: ModelProfile
+    model: Any = None
+    tokenizer: Any = None
+    head: Any = None # Extension for DINO head architecture
+    stats: dict = None # Extension for anomaly scoring
+    is_loaded: bool = False
+    is_loading: bool = False
+    load_error: str | None = None
+    load_time_ms: int = 0
+    last_used: datetime | None = None
+    current_device: str = "unloaded"
+    
+    @property
+    def is_available(self) -> bool:
+        return self.is_loaded and self.load_error is None
+
+class ModelRegistry:
+    def __init__(self):
+        self._states: dict[str, ModelState] = {
+            name: ModelState(profile=profile) for name, profile in MODEL_PROFILES.items()
+        }
+        self._locks: dict[str, asyncio.Lock] = {
+            name: asyncio.Lock() for name in MODEL_PROFILES
+        }
+        self._gpu_budget_mb = settings.GPU_VRAM_BUDGET_MB
+    
+    async def startup_load(self):
+        ml_logger.logger.info("Starting model registry startup")
+        sorted_models = sorted(MODEL_PROFILES.values(), key=lambda m: m.load_priority)
+        
+        for profile in sorted_models:
+            if profile.device_preference == "cpu":
+                await self._load_model(profile.name)
+            else:
+                if self._get_used_vram() + profile.vram_mb <= self._gpu_budget_mb:
+                    await self._load_model(profile.name)
+                else:
+                    ml_logger.logger.warning(f"Skipping GPU load for {profile.name}: VRAM budget exceeded. Will load on CPU on first request.")
+        
+        loaded = [n for n, s in self._states.items() if s.is_available]
+        failed = [n for n, s in self._states.items() if s.load_error]
+        required_failed = [n for n in failed if MODEL_PROFILES[n].is_required]
+        
+        if required_failed:
+            raise RuntimeError(f"Critical models failed to load: {required_failed}. Check logs.")
+            
+        ml_logger.logger.info("Registry startup complete", extra={"loaded": loaded, "failed": failed, "vram_used_mb": self._get_used_vram()})
+
+    async def get(self, model_name: str) -> ModelState:
+        if model_name not in self._states:
+            raise ValueError(f"Unknown model: {model_name}")
+        
+        state = self._states[model_name]
+        if not state.is_available and not state.is_loading:
+            await self._load_model(model_name)
+            
+        self._states[model_name].last_used = datetime.now(UTC)
+        return self._states[model_name]
+
+    def is_available(self, model_name: str) -> bool:
+        return self._states.get(model_name, ModelState(ModelProfile("","","","cpu",0,0,0,False))).is_available
+
+    async def _load_model(self, model_name: str):
+        async with self._locks[model_name]:
+            state = self._states[model_name]
+            if state.is_available:
+                return
+            
+            state.is_loading = True
+            start_time = time.monotonic()
+            
+            try:
+                profile = state.profile
+                device = self._resolve_device(profile)
+                
+                if device == "cuda":
+                    needed = profile.vram_mb
+                    available = self._gpu_budget_mb - self._get_used_vram()
+                    if available < needed:
+                        evicted = await self._evict_lru_gpu_model(except_model=model_name)
+                        if evicted:
+                            ml_logger.logger.info(f"Evicted {evicted} to make room for {model_name}")
+                
+                # Fetch objects securely
+                result = await asyncio.to_thread(self._load_model_sync, model_name, profile, device)
+                
+                load_time_ms = int((time.monotonic() - start_time) * 1000)
+                
+                state.model = result.get('model')
+                state.tokenizer = result.get('tokenizer')
+                state.head = result.get('head')
+                state.stats = result.get('stats')
+                
+                state.is_loaded = True
+                state.load_error = None
+                state.load_time_ms = load_time_ms
+                state.current_device = device
+                
+                ml_logger.log_model_load(model_name, device, load_time_ms, vram_delta_mb=profile.vram_mb if device == "cuda" else None)
+                
+            except Exception as e:
+                state.load_error = str(e)
+                state.is_loaded = False
+                ml_logger.logger.error(f"Failed to load model {model_name}: {e}", exc_info=True)
+                if MODEL_PROFILES[model_name].is_required:
+                    raise
+            finally:
+                state.is_loading = False
+
+    def _load_model_sync(self, name: str, profile: ModelProfile, device: str) -> dict:
+        cache_dir = settings.MODEL_CACHE_DIR / profile.local_cache_subdir
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        
+        if name == "dino_anomaly":
+            from transformers import AutoImageProcessor, AutoModel as ViTModel
+            processor = AutoImageProcessor.from_pretrained(profile.hf_model_id, cache_dir=cache_dir)
+            model = ViTModel.from_pretrained(profile.hf_model_id, cache_dir=cache_dir).to(device)
+            model.eval()
+            
+            # Simulated Projection Head Loading Logic
+            head = None 
+            stats = {"mean": 0.001, "std": 0.0005} 
+            head_path = settings.MODEL_CACHE_DIR / "anomaly_head.pt"
+            stats_path = settings.MODEL_CACHE_DIR / "anomaly_stats.json"
+            
+            if head_path.exists():
+                 # We will define DINOProjectionHead in the vision module later
+                 import json
+                 pass 
+            else:
+                 logger.warning("Using untrained head — anomaly scores may be unreliable")
+            
+            if stats_path.exists():
+                import json
+                stats = json.loads(stats_path.read_text())
+
+            return {"model": model, "tokenizer": processor, "head": head, "stats": stats}
+            
+        elif name == "biobert_ner":
+            fine_tuned_path = settings.MODEL_CACHE_DIR / "biobert_ner_finetuned"
+            model_path = str(fine_tuned_path) if fine_tuned_path.exists() else profile.hf_model_id
+            
+            from transformers import AutoModelForTokenClassification, AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained(model_path, cache_dir=cache_dir)
+            model = AutoModelForTokenClassification.from_pretrained(model_path, cache_dir=cache_dir).to(device)
+            model.eval()
+            return {"model": model, "tokenizer": tokenizer}
+            
+        elif name == "whisper_tiny":
+            import whisper
+            model = whisper.load_model("tiny", device="cpu", download_root=str(cache_dir))
+            return {"model": model, "tokenizer": None}
+            
+        elif name == "biogpt_base":
+            from transformers import BioGptForCausalLM, BioGptTokenizer
+            tokenizer = BioGptTokenizer.from_pretrained(profile.hf_model_id, cache_dir=cache_dir)
+            model = BioGptForCausalLM.from_pretrained(profile.hf_model_id, cache_dir=cache_dir)
+            model.eval()
+            return {"model": model, "tokenizer": tokenizer}
+            
+        elif name == "minilm":
+            from sentence_transformers import SentenceTransformer
+            model = SentenceTransformer(profile.hf_model_id, cache_folder=str(cache_dir))
+            return {"model": model, "tokenizer": None}
+            
+        elif name == "biomedvlp":
+            from transformers import AutoModel, AutoTokenizer
+            free_vram = torch.cuda.mem_get_info()[0] // 1024 // 1024 if torch.cuda.is_available() else 0
+            if free_vram < 950 and device == "cuda":
+                device = "cpu"
+                logger.warning("Insufficient VRAM for BiomedVLP — loading on CPU (slower)")
+                
+            # Trust remote code is required for custom Microsoft implementation
+            tokenizer = AutoTokenizer.from_pretrained(profile.hf_model_id, cache_dir=cache_dir, trust_remote_code=True)
+            model = AutoModel.from_pretrained(profile.hf_model_id, cache_dir=cache_dir, trust_remote_code=True).to(device)
+            model.eval()
+            return {"model": model, "tokenizer": tokenizer}
+            
+        else:
+            raise ValueError(f"No loader defined for model: {name}")
+
+    def _resolve_device(self, profile: ModelProfile) -> str:
+        if profile.device_preference == "cpu":
+            return "cpu"
+        if profile.device_preference == "cuda":
+            if not torch.cuda.is_available():
+                ml_logger.logger.warning(f"CUDA not available, loading {profile.name} on CPU")
+                return "cpu"
+            return "cuda"
+        if profile.device_preference == "auto":
+            if torch.cuda.is_available():
+                free_vram = self._gpu_budget_mb - self._get_used_vram()
+                if free_vram >= profile.vram_mb:
+                    return "cuda"
+            return "cpu"
+
+    def _get_used_vram(self) -> int:
+        return sum(s.profile.vram_mb for s in self._states.values() if s.is_available and s.current_device == "cuda")
+
+    async def _evict_lru_gpu_model(self, except_model: str) -> str | None:
+        gpu_models = [
+            (name, state) for name, state in self._states.items()
+            if state.is_available and state.current_device == "cuda" and name != except_model
+        ]
+        if not gpu_models:
+            return None
+            
+        lru_name, _ = min(gpu_models, key=lambda x: x[1].last_used or datetime.min.replace(tzinfo=UTC))
+        await asyncio.to_thread(self._move_to_cpu, lru_name)
+        return lru_name
+
+    def _move_to_cpu(self, model_name: str):
+        state = self._states[model_name]
+        if state.model is not None and hasattr(state.model, "cpu"):
+            state.model = state.model.cpu()
+            torch.cuda.empty_cache()
+            state.current_device = "cpu"
+            ml_logger.logger.info(f"Moved {model_name} to CPU")
+
+    def get_status(self) -> dict:
+        return {
+            "models": {
+                name: {
+                    "is_available": state.is_available,
+                    "device": state.current_device,
+                    "load_error": state.load_error,
+                    "load_time_ms": state.load_time_ms,
+                    "last_used": state.last_used.isoformat() if state.last_used else None,
+                    "vram_mb": state.profile.vram_mb if state.current_device == "cuda" else 0
+                }
+                for name, state in self._states.items()
+            },
+            "gpu_budget_mb": self._gpu_budget_mb,
+            "gpu_used_mb": self._get_used_vram(),
+            "gpu_free_mb": self._gpu_budget_mb - self._get_used_vram()
+        }
+
+model_registry = ModelRegistry()

backend/orchestration/pipeline.py

@@ -0,0 +1,184 @@
+import time
+import asyncio
+import logging
+import torch
+from pathlib import Path
+from datetime import datetime, UTC
+
+from backend.ml.registry import ModelRegistry
+from backend.core.exceptions import InvalidFileError, InferenceError, ModelNotLoadedError
+from backend.api.v1.schemas.analysis import (
+    AnalysisResult, VisionResult, NLPResult, FusionResult, ProcessingTimings
+)
+
+logger = logging.getLogger(__name__)
+
+class AnalysisPipeline:
+    def __init__(self, registry: ModelRegistry):
+        self.registry = registry
+
+    async def run(self, session_id: str, image_path: Path, symptoms_text: str) -> AnalysisResult:
+        timings = {}
+        warnings = []
+        vision_result = None
+        nlp_result = None
+        fusion_result = None
+        report_text = None
+        
+        # ── STEP 1: VALIDATE INPUT ────────────────────────────
+        t0 = time.monotonic()
+        try:
+            if not image_path.exists():
+                raise InvalidFileError("Image file not found")
+            processed_image_path = await asyncio.to_thread(self._preprocess_image, image_path)
+        except Exception as e:
+            raise InferenceError(f"Input validation failed: {e}")
+        timings["preprocess_ms"] = int((time.monotonic() - t0) * 1000)
+
+        # ── STEP 2: VISION ANALYSIS (GPU) ─────────────────────
+        t0 = time.monotonic()
+        try:
+            # We wrap this in resilience layers in resilience.py
+            vision_result = await self._run_vision(processed_image_path)
+        except Exception as e:
+            logger.warning(f"Vision analysis failed for session {session_id}: {e}")
+            warnings.append(f"Vision analysis unavailable: {type(e).__name__}")
+        timings["vision_ms"] = int((time.monotonic() - t0) * 1000)
+
+        # ── STEP 3: VRAM CLEANUP ───────────────────────────────
+        if vision_result is not None:
+            if torch.cuda.is_available():
+                await asyncio.to_thread(torch.cuda.empty_cache)
+                await asyncio.sleep(0.1)
+
+        # ── STEP 4: NLP ANALYSIS (GPU/CPU) ────────────────────
+        t0 = time.monotonic()
+        if symptoms_text.strip():
+            try:
+                nlp_result = await self._run_nlp(symptoms_text)
+            except Exception as e:
+                logger.warning(f"NLP analysis failed for session {session_id}: {e}")
+                warnings.append(f"NLP analysis unavailable: {type(e).__name__}")
+        else:
+            warnings.append("No symptoms text provided — NLP analysis skipped")
+        timings["nlp_ms"] = int((time.monotonic() - t0) * 1000)
+
+        # ── STEP 5: MULTIMODAL FUSION ─────────────────────────
+        t0 = time.monotonic()
+        if vision_result is not None and nlp_result is not None:
+            try:
+                fusion_result = await self._run_fusion(processed_image_path, symptoms_text)
+            except Exception as e:
+                logger.warning(f"Fusion failed for session {session_id}: {e}")
+                warnings.append(f"Multimodal fusion unavailable: {type(e).__name__}")
+        else:
+            warnings.append("Fusion skipped: requires both vision and NLP results")
+        timings["fusion_ms"] = int((time.monotonic() - t0) * 1000)
+
+        # ── STEP 6: REPORT GENERATION ─────────────────────────
+        t0 = time.monotonic()
+        try:
+            report_text = await self._generate_report(vision_result, nlp_result, fusion_result)
+        except Exception as e:
+            logger.warning(f"Report generation failed: {e}")
+            report_text = self._fallback_report(vision_result, nlp_result)
+            warnings.append("Using template report — AI report generation unavailable")
+        timings["report_ms"] = int((time.monotonic() - t0) * 1000)
+
+        # ── STEP 7: DETERMINE OVERALL STATUS ──────────────────
+        if vision_result is None and nlp_result is None:
+            overall_status = "FAILED"
+        elif vision_result is None or nlp_result is None:
+            overall_status = "PARTIAL"
+        else:
+            overall_status = "COMPLETE"
+
+        timings["total_ms"] = sum(timings.values())
+
+        return AnalysisResult(
+            session_id=session_id, patient_id="", timestamp=datetime.now(UTC),
+            vision=vision_result, nlp=nlp_result, fusion=fusion_result,
+            report_text=report_text, overall_status=overall_status,
+            timings=ProcessingTimings(**timings), warnings=warnings
+        )
+
+    def _preprocess_image(self, image_path: Path) -> Path:
+        from PIL import Image
+        with Image.open(image_path) as img:
+            img = img.convert("RGB")
+            img = img.resize((224, 224), Image.LANCZOS)
+            output_path = image_path.with_suffix(".processed.png")
+            img.save(output_path, "PNG")
+        return output_path
+
+    async def _run_vision(self, image_path: Path) -> VisionResult:
+        from backend.ml.vision.anomaly import AnomalyDetector
+        from backend.ml.vision.gradcam import GradCAM
+        
+        state = await self.registry.get("dino_anomaly")
+        if not state.is_available:
+            raise ModelNotLoadedError("Vision model unavailable")
+            
+        anomaly_score, model_confidence = await asyncio.to_thread(
+            AnomalyDetector.score, image_path, state.model, state.head, state.stats, state.current_device
+        )
+        
+        heatmap_b64, top_regions = await asyncio.to_thread(
+            GradCAM.generate, image_path, state.model, anomaly_score
+        )
+        
+        risk_level = "LOW" if anomaly_score < 40 else "MEDIUM" if anomaly_score < 70 else "HIGH"
+        return VisionResult(
+            anomaly_score=round(anomaly_score, 1), risk_level=risk_level,
+            heatmap_base64=heatmap_b64, top_regions=top_regions, model_confidence=model_confidence
+        )
+
+    async def _run_nlp(self, text: str) -> NLPResult:
+        from backend.ml.nlp.ner import NERExtractor
+        from backend.ml.nlp.classifier import DiseaseClassifier
+        
+        ner_state = await self.registry.get("biobert_ner")
+        entities = await asyncio.to_thread(
+            NERExtractor.extract, text, ner_state.model, ner_state.tokenizer, ner_state.current_device
+        )
+        
+        diagnosis = await asyncio.to_thread(DiseaseClassifier.classify, text, entities)
+        return NLPResult(
+            entities=entities, primary_diagnosis=diagnosis["primary"],
+            diagnosis_confidence=diagnosis["confidence"], differential=diagnosis["differential"]
+        )
+
+    async def _run_fusion(self, image_path: Path, text: str) -> FusionResult:
+        from backend.ml.fusion.medclip import MultimodalFusion
+        
+        state = await self.registry.get("biomedvlp")
+        if not state.is_available:
+            raise ModelNotLoadedError("Fusion model unavailable")
+            
+        similarity, alignment = await asyncio.to_thread(
+            MultimodalFusion.compute_similarity, image_path, text, state.model, state.tokenizer, state.current_device
+        )
+        final_risk = "HIGH" if similarity < 0.3 else "MEDIUM" if similarity < 0.7 else "LOW"
+        return FusionResult(image_text_similarity=round(similarity, 3), alignment=alignment, final_risk=final_risk)
+
+    async def _generate_report(self, vision: VisionResult | None, nlp: NLPResult | None, fusion: FusionResult | None) -> str:
+        from backend.ml.rag.generator import ReportGenerator
+        
+        state = await self.registry.get("biogpt_base")
+        if not state.is_available:
+            raise ModelNotLoadedError("Report generation unavailable")
+            
+        return await asyncio.to_thread(
+            ReportGenerator.generate, vision, nlp, fusion, state.model, state.tokenizer
+        )
+
+    def _fallback_report(self, vision: VisionResult | None, nlp: NLPResult | None) -> str:
+        parts = ["## AI-Assisted Analysis Report\n\n*Note: This is an automated template report.*\n"]
+        if vision:
+            parts.append(f"**Imaging Findings:** Anomaly score of {vision.anomaly_score}/100 indicates {vision.risk_level.lower()} risk findings.")
+        if nlp:
+            diseases = ", ".join(nlp.entities.diseases) if nlp.entities.diseases else "none identified"
+            symptoms = ", ".join(nlp.entities.symptoms) if nlp.entities.symptoms else "none documented"
+            parts.append(f"**Clinical Impression:** {nlp.primary_diagnosis} (confidence: {nlp.diagnosis_confidence:.0%}). Identified conditions: {diseases}. Symptoms: {symptoms}.")
+        parts.append("\n**Recommendation:** Please consult a licensed physician for diagnosis and treatment.")
+        return "\n".join(parts)

backend/orchestration/queue.py

@@ -0,0 +1,198 @@
+import asyncio
+import uuid
+import logging
+from pathlib import Path
+from datetime import datetime, UTC
+from sqlalchemy import select, func
+
+from backend.db.models import AnalysisTask, AnalysisSession
+from backend.core.exceptions import TaskNotFoundError, SessionAccessDeniedError
+from backend.core.logging_config import ml_logger
+
+logger = logging.getLogger(__name__)
+
+class AnalysisTaskQueue:
+    MAX_CONCURRENT = 2
+    WORKER_SLEEP_SECONDS = 5
+    
+    def __init__(self, db_session_factory, pipeline):
+        self._db_factory = db_session_factory
+        self._pipeline = pipeline
+        self._new_task_event = asyncio.Event()
+        self._active_count = 0
+        self._active_lock = asyncio.Lock()
+        self._worker_task: asyncio.Task | None = None
+        self._is_running = False
+    
+    async def start(self):
+        self._is_running = True
+        self._worker_task = asyncio.create_task(self._worker_loop())
+        logger.info("Task queue worker started")
+    
+    async def stop(self):
+        logger.info("Task queue stopping...")
+        self._is_running = False
+        self._new_task_event.set()
+        
+        deadline = asyncio.get_event_loop().time() + 60
+        while self._active_count > 0:
+            if asyncio.get_event_loop().time() > deadline:
+                logger.warning(f"Shutdown timeout: {self._active_count} tasks still running")
+                break
+            await asyncio.sleep(1)
+        
+        if self._worker_task:
+            self._worker_task.cancel()
+
+    async def submit(self, session_id: str, user_id: str | None, image_path: str, symptoms_text: str, priority: int = 1) -> str:
+        task_id = str(uuid.uuid4())
+        async with self._db_factory() as db:
+            task = AnalysisTask(
+                id=task_id, session_id=session_id, user_id=user_id,
+                status="PENDING", priority=priority,
+                image_path=image_path, symptoms_text=symptoms_text
+            )
+            db.add(task)
+            await db.commit()
+            
+        self._new_task_event.set()
+        return task_id
+
+    async def get_status(self, task_id: str) -> dict:
+        async with self._db_factory() as db:
+            task = await db.get(AnalysisTask, task_id)
+            if not task:
+                raise TaskNotFoundError()
+                
+            position = None
+            if task.status == "PENDING":
+                result = await db.execute(
+                    select(func.count(AnalysisTask.id)).where(
+                        AnalysisTask.status == "PENDING",
+                        AnalysisTask.priority >= task.priority,
+                        AnalysisTask.created_at < task.created_at
+                    )
+                )
+                position = result.scalar_one() + 1
+                
+            estimated_wait = None
+            if position:
+                slots_until_ours = max(0, position - self.MAX_CONCURRENT)
+                estimated_wait = slots_until_ours * 45
+                
+            return {
+                "task_id": task_id, "session_id": task.session_id, "status": task.status,
+                "position_in_queue": position, "estimated_wait_seconds": estimated_wait,
+                "started_at": task.started_at, "completed_at": task.completed_at,
+                "error_message": task.error_message
+            }
+
+    async def cancel(self, task_id: str, user_id: str) -> bool:
+        async with self._db_factory() as db:
+            task = await db.get(AnalysisTask, task_id)
+            if not task:
+                raise TaskNotFoundError()
+            if task.user_id != user_id:
+                raise SessionAccessDeniedError()
+            if task.status != "PENDING":
+                return False
+                
+            task.status = "CANCELLED"
+            image_path = Path(task.image_path)
+            if image_path.exists():
+                image_path.unlink()
+                
+            await db.commit()
+            return True
+
+    async def _worker_loop(self):
+        logger.info("Worker loop started")
+        while self._is_running:
+            try:
+                await asyncio.wait_for(
+                    asyncio.shield(self._new_task_event.wait()), timeout=self.WORKER_SLEEP_SECONDS
+                )
+                self._new_task_event.clear()
+            except asyncio.TimeoutError:
+                pass
+                
+            if not self._is_running:
+                break
+                
+            while self._active_count < self.MAX_CONCURRENT:
+                task = await self._fetch_next_task()
+                if not task:
+                    break
+                asyncio.create_task(self._process_task(task))
+                
+        logger.info("Worker loop stopped")
+
+    async def _fetch_next_task(self) -> AnalysisTask | None:
+        async with self._db_factory() as db:
+            result = await db.execute(
+                select(AnalysisTask).where(AnalysisTask.status == "PENDING")
+                .order_by(AnalysisTask.priority.desc(), AnalysisTask.created_at.asc())
+                .limit(1)
+                .with_for_update(skip_locked=True)
+            )
+            return result.scalar_one_or_none()
+
+    async def _process_task(self, task: AnalysisTask):
+        async with self._active_lock:
+            self._active_count += 1
+            
+        try:
+            async with self._db_factory() as db:
+                task.status = "PROCESSING"
+                task.started_at = datetime.now(UTC)
+                await db.merge(task)
+                await db.commit()
+                
+            result = await self._pipeline.run(
+                session_id=task.session_id,
+                image_path=Path(task.image_path),
+                symptoms_text=task.symptoms_text or ""
+            )
+            
+            async with self._db_factory() as db:
+                db_task = await db.get(AnalysisTask, task.id)
+                db_task.status = "COMPLETED"
+                db_task.completed_at = datetime.now(UTC)
+                
+                session = await db.get(AnalysisSession, task.session_id)
+                session.result_json = result.model_dump()
+                session.status = "READY"
+                session.risk_level = result.vision.risk_level if result.vision else "UNKNOWN"
+                
+                await db.commit()
+                
+            ml_logger.log_pipeline_step(
+                "full_pipeline", "COMPLETED", 
+                int((datetime.now(UTC) - task.started_at).total_seconds() * 1000), task.session_id
+            )
+            
+        except Exception as e:
+            logger.error(f"Task {task.id} failed: {e}", exc_info=True)
+            async with self._db_factory() as db:
+                db_task = await db.get(AnalysisTask, task.id)
+                db_task.attempt_count = (db_task.attempt_count or 0) + 1
+                
+                if db_task.attempt_count < 3:
+                    db_task.status = "PENDING"
+                    self._new_task_event.set()
+                else:
+                    db_task.status = "FAILED"
+                    db_task.error_message = str(e)[:500]
+                    session = await db.get(AnalysisSession, task.session_id)
+                    session.status = "FAILED"
+                    session.error_message = "Analysis failed after 3 attempts"
+                await db.commit()
+                
+        finally:
+            try:
+                Path(task.image_path).unlink(missing_ok=True)
+            except Exception:
+                pass
+            async with self._active_lock:
+                self._active_count -= 1
+            self._new_task_event.set()

frontend/app/(auth)/callback/page.jsx

@@ -0,0 +1,59 @@
+'use client'
+import { useEffect, useState } from 'react'
+import { useRouter, useSearchParams } from 'next/navigation'
+import { useAuth } from '@/lib/auth/AuthContext'
+import Link from 'next/link'
+
+export default function AuthCallbackPage() {
+    const searchParams = useSearchParams()
+    const router = useRouter()
+    const { loginWithToken } = useAuth()
+    const [errorMsg, setErrorMsg] = useState(null)
+
+    useEffect(() => {
+        const processToken = async () => {
+            const token = searchParams.get('token')
+            const error = searchParams.get('error')
+
+            if (error) {
+                setErrorMsg(`Authentication failed: ${error}`)
+                return
+            }
+
+            if (token) {
+                try {
+                    await loginWithToken(token)
+                    // Clean URL and redirect
+                    router.replace('/upload')
+                } catch (err) {
+                    setErrorMsg("Failed to establish secure session.")
+                }
+            } else {
+                setErrorMsg("No authentication token provided.")
+            }
+        }
+
+        processToken()
+    }, [searchParams, loginWithToken, router])
+
+    if (errorMsg) {
+        return (
+            <div className="flex flex-col items-center justify-center min-h-screen p-4 text-center">
+                <div className="p-6 bg-red-50 border border-red-200 rounded-lg shadow-sm">
+                    <h2 className="text-lg font-semibold text-red-700 mb-2">Error</h2>
+                    <p className="text-red-600 mb-4">{errorMsg}</p>
+                    <Link href="/login" className="px-4 py-2 text-white bg-blue-600 rounded hover:bg-blue-700 transition">
+                        Return to Login
+                    </Link>
+                </div>
+            </div>
+        )
+    }
+
+    return (
+        <div className="flex flex-col items-center justify-center min-h-screen">
+            <div className="w-8 h-8 border-4 border-blue-600 border-t-transparent rounded-full animate-spin mb-4" />
+            <p className="text-gray-600 font-medium">Securing session...</p>
+        </div>
+    )
+}

frontend/components/auth/ProtectedRoute.jsx

@@ -0,0 +1,23 @@
+'use client'
+import { useEffect } from 'react'
+import { useRouter, usePathname } from 'next/navigation'
+import { useAuth } from '@/lib/auth/AuthContext'
+
+export default function ProtectedRoute({ children }) {
+    const { isAuthenticated, isLoading } = useAuth()
+    const router = useRouter()
+    const pathname = usePathname()
+
+    useEffect(() => {
+        if (!isLoading && !isAuthenticated) {
+            sessionStorage.setItem('intendedPath', pathname)
+            router.push('/login')
+        }
+    }, [isLoading, isAuthenticated, router, pathname])
+
+    if (isLoading) {
+        return <div className="flex items-center justify-center min-h-screen">Loading secure environment...</div>
+    }
+
+    return isAuthenticated ? children : null
+}

frontend/components/auth/PublicOnlyRoute.jsx

@@ -0,0 +1,19 @@
+'use client'
+import { useEffect } from 'react'
+import { useRouter } from 'next/navigation'
+import { useAuth } from '@/lib/auth/AuthContext'
+
+export default function PublicOnlyRoute({ children }) {
+    const { isAuthenticated, isLoading } = useAuth()
+    const router = useRouter()
+
+    useEffect(() => {
+        if (!isLoading && isAuthenticated) {
+            router.push('/upload')
+        }
+    }, [isLoading, isAuthenticated, router])
+
+    if (isLoading) return null
+
+    return !isAuthenticated ? children : null
+}

frontend/next.config.js

@@ -0,0 +1,43 @@
+/** @type {import('next').NextConfig} */
+
+const securityHeaders = [
+  {
+    key: "Content-Security-Policy",
+    value: [
+      "default-src 'self'",
+      "script-src 'self' 'unsafe-eval' 'unsafe-inline'",  
+      "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com",
+      "font-src 'self' https://fonts.gstatic.com",
+      "img-src 'self' data: blob: https:", 
+      `connect-src 'self' ${process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8000'}`,
+      "frame-src 'none'",
+      "object-src 'none'",
+      "base-uri 'self'",
+      "form-action 'self'",
+    ].join("; ")
+  },
+  { key: "X-DNS-Prefetch-Control", value: "on" },
+  { key: "X-Frame-Options", value: "SAMEORIGIN" },
+  { key: "X-Content-Type-Options", value: "nosniff" },
+  { key: "Referrer-Policy", value: "strict-origin-when-cross-origin" },
+];
+
+const nextConfig = {
+  output: "standalone",
+  experimental: {
+    serverActions: true,
+  },
+  images: {
+    domains: [(process.env.NEXT_PUBLIC_API_URL || 'localhost').replace(/^https?:\/\//, '').split(':')[0]],
+  },
+  async headers() {
+    return [
+      {
+        source: "/(.*)",
+        headers: securityHeaders,
+      },
+    ];
+  },
+};
+
+module.exports = nextConfig;

scripts/security_audit.py

@@ -0,0 +1,47 @@
+import sys
+import httpx
+import asyncio
+
+API_URL = "http://localhost:8000"
+
+async def run_audit():
+    print("\n--- MedSight AI Security Audit ---\n")
+    passed = 0
+    total = 0
+
+    def check(name, condition, error_msg):
+        nonlocal passed, total
+        total += 1
+        if condition:
+            print(f"✅ PASS: {name}")
+            passed += 1
+        else:
+            print(f"❌ FAIL: {name} - {error_msg}")
+
+    try:
+        async with httpx.AsyncClient() as client:
+            # 1. Test Health / Headers
+            res = await client.get(f"{API_URL}/api/v1/health")
+            headers = res.headers
+            check("X-Content-Type-Options", headers.get("x-content-type-options") == "nosniff", "Header missing or incorrect")
+            check("X-Frame-Options", headers.get("x-frame-options") == "DENY", "Header missing or incorrect")
+            check("Permissions-Policy", "microphone=(self)" in headers.get("permissions-policy", ""), "Microphone permission misconfigured")
+            check("No Server Header", "server" not in headers, "Server header is leaking framework info")
+
+            # 2. Test CORS Rejection
+            cors_res = await client.options(
+                f"{API_URL}/api/v1/health",
+                headers={"Origin": "http://evil-domain.com", "Access-Control-Request-Method": "GET"}
+            )
+            # Depending on FastAPI config, it might strip CORS headers or return 400. 
+            # The key is it shouldn't return Access-Control-Allow-Origin: http://evil-domain.com
+            check("CORS Restrictions", cors_res.headers.get("access-control-allow-origin") != "http://evil-domain.com", "CORS allowed untrusted origin")
+
+    except httpx.ConnectError:
+        print("❌ Server not running. Start it with: make run-dev")
+        sys.exit(1)
+
+    print(f"\nAudit Complete: {passed}/{total} Passed\n")
+
+if __name__ == "__main__":
+    asyncio.run(run_audit())