fix: harden credentials, sessions, uploads, and retrieval

api/bigrag/models/chat_readiness.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class ChatReadinessResponse(BaseModel):
+    credential_ready: bool
+    credential_source: Literal["saved", "instance"] | None = None
+    reason: (
+        Literal["missing_credentials", "credential_blocked", "decryption_unavailable"] | None
+    ) = None

api/bigrag/routers/admin_api_keys.py

@@ -166,7 +166,7 @@ async def update_api_key(
     if body.active is not None:
         key.active = body.active
         fields.append("active")
-    if body.expires_at is not None:
+    if "expires_at" in body.model_fields_set:
         key.expires_at = body.expires_at
         fields.append("expires_at")
     existing = dict(key.permissions or {})

api/bigrag/routers/auth.py

@@ -276,6 +276,7 @@ async def whoami(user: dict = Depends(get_current_user)) -> WhoamiResponse:
 async def change_password(
     body: ChangePasswordRequest,
     request: Request,
+    response: Response,
     user: dict = Depends(require_session),
     session: AsyncSession = Depends(get_session),
 ) -> StatusResponse:
@@ -289,6 +290,7 @@ async def change_password(
     await session.execute(sa.delete(UserSession).where(UserSession.user_id == target.id))
     await session.commit()
     await invalidate_auth_principals()
+    await clear_session_cookie(response)
     audit.record(
         request,
         user=user,

api/bigrag/routers/documents/crud.py

@@ -59,12 +59,6 @@ async def upload_document(
     upload_limits = await get_values(["max_upload_size_mb"])
     max_upload_size_mb = upload_limits["max_upload_size_mb"]
     max_size = max_upload_size_mb * 1024 * 1024
-    content_length = request.headers.get("content-length")
-    if content_length and int(content_length) > max_size:
-        raise HTTPException(
-            status_code=413,
-            detail=f"File too large. Max size: {max_upload_size_mb}MB",
-        )
 
     tmp_path, content_hash, file_size = await validated_upload_to_temp(
         file, file_ext, max_size=max_size

api/bigrag/routers/documents_batch.py

@@ -121,12 +121,6 @@ async def batch_upload_documents(
     max_batch_upload_size_mb = upload_limits["max_batch_upload_size_mb"]
     max_size = max_upload_size_mb * 1024 * 1024
     batch_max_size = max_batch_upload_size_mb * 1024 * 1024
-    content_length = request.headers.get("content-length")
-    if content_length and int(content_length) > batch_max_size:
-        raise HTTPException(
-            status_code=413,
-            detail=f"Batch upload too large. Max size: {max_batch_upload_size_mb}MB",
-        )
     budget = UploadBudget(batch_max_size)
     shared_meta = metadata_or_400(
         collection,

api/bigrag/routers/evaluation.py

@@ -68,10 +68,11 @@ def _ndcg_at_k(hit_ids: list[str], expected: set[str]) -> float:
     if not expected:
         return 0.0
     dcg = 0.0
+    credited: set[str] = set()
     for idx, h in enumerate(hit_ids, start=1):
-        rel = 1 if h in expected else 0
-        if rel:
+        if h in expected and h not in credited:
             dcg += 1.0 / math.log2(idx + 1)
+            credited.add(h)
     k = len(hit_ids)
     ideal = sum(1.0 / math.log2(i + 1) for i in range(1, min(len(expected), k) + 1))
     return dcg / ideal if ideal else 0.0

api/bigrag/routers/preferences.py

@@ -3,14 +3,16 @@
 import uuid
 
 import sqlalchemy as sa
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Response
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from bigrag.db.models import UserPreference
 from bigrag.db.session import get_session
 from bigrag.logging import get_logger
 from bigrag.middleware.auth import require_session
+from bigrag.models.chat_readiness import ChatReadinessResponse
+from bigrag.services.chat.readiness import get_chat_readiness
 from bigrag.services.preferences import (
     _deep_merge,
     _encrypt_sensitive,
@@ -25,6 +27,16 @@
 router = APIRouter(prefix="/v1/auth/preferences", tags=["auth"])
 
 
+@router.get("/chat-readiness", response_model=ChatReadinessResponse)
+async def chat_readiness(
+    response: Response,
+    user: dict = Depends(require_session),
+    session: AsyncSession = Depends(get_session),
+) -> ChatReadinessResponse:
+    response.headers["Cache-Control"] = "private, no-store"
+    return await get_chat_readiness(session, user)
+
+
 @router.get("", response_model=dict[str, dict])
 async def get_preferences(
     user: dict = Depends(require_session),
@@ -53,7 +65,7 @@ async def update_preferences(
     )
     existing = dict(existing_row.data) if existing_row else {}
     incoming = _normalize_sensitive(incoming)
-    await _validate_sensitive(incoming)
+    incoming = await _validate_sensitive(incoming)
     incoming = _encrypt_sensitive(incoming)
     merged = _remove_cleared_sensitive(_deep_merge(existing, incoming), incoming)
 

api/bigrag/routers/webhooks/crud.py

@@ -142,7 +142,7 @@ async def update_webhook(
     if body.events is not None:
         wh.events = body.events
         fields.append("events")
-    if body.collections is not None:
+    if "collections" in body.model_fields_set:
         wh.collections = body.collections
         fields.append("collections")
     if body.active is not None:

api/bigrag/services/analytics.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import asyncio
-
 import sqlalchemy as sa
 
 from bigrag.db.engine import session_factory
@@ -49,11 +47,9 @@ async def collection_analytics(collection_name: str) -> dict:
             )
         ).all()
 
-        stats_24h, stats_7d, stats_30d = await asyncio.gather(
-            _period_stats(session, collection_name, 1),
-            _period_stats(session, collection_name, 7),
-            _period_stats(session, collection_name, 30),
-        )
+        stats_24h = await _period_stats(session, collection_name, 1)
+        stats_7d = await _period_stats(session, collection_name, 7)
+        stats_30d = await _period_stats(session, collection_name, 30)
 
     result = {
         "collection": collection_name,

api/bigrag/services/chat/completion.py

@@ -2,8 +2,10 @@
 
 import asyncio
 from collections.abc import AsyncIterator
+from contextlib import suppress
 from datetime import UTC, datetime
 
+import anyio
 from fastapi import Request
 
 from bigrag.db.engine import session_factory
@@ -87,28 +89,38 @@ async def stream_chat_completion(
 
         content_parts: list[str] = []
         stream = _stream_model(prepared).__aiter__()
+        pending_delta: asyncio.Task[str] | None = None
         try:
             while True:
                 if request is not None and await request.is_disconnected():
                     return
-                try:
-                    delta = await asyncio.wait_for(
-                        stream.__anext__(), timeout=_HEARTBEAT_INTERVAL_SECONDS
-                    )
-                except TimeoutError:
+                if pending_delta is None:
+                    pending_delta = asyncio.create_task(stream.__anext__())
+                completed, _ = await asyncio.wait(
+                    {pending_delta}, timeout=_HEARTBEAT_INTERVAL_SECONDS
+                )
+                if not completed:
                     yield _HEARTBEAT
                     continue
+                try:
+                    delta = pending_delta.result()
                 except StopAsyncIteration:
                     break
+                pending_delta = None
                 content_parts.append(delta)
                 yield _sse("delta", {"delta": delta})
         finally:
-            aclose = getattr(stream, "aclose", None)
-            if aclose is not None:
-                try:
-                    await aclose()
-                except Exception:
-                    pass
+            with anyio.CancelScope(shield=True):
+                if pending_delta is not None:
+                    pending_delta.cancel()
+                    with suppress(asyncio.CancelledError, Exception):
+                        await pending_delta
+                aclose = getattr(stream, "aclose", None)
+                if aclose is not None:
+                    try:
+                        await aclose()
+                    except Exception:
+                        pass
 
         assistant = _chat_message_response(
             id=str(uuid7()),

api/bigrag/services/chat/provider.py

@@ -99,13 +99,14 @@ async def _stream_model(prepared: PreparedChatTurn) -> AsyncIterator[str]:
                 ),
                 timeout=_MODEL_TIMEOUT_SECONDS,
             )
-            async for chunk in stream:
-                choices = getattr(chunk, "choices", None) or []
-                if not choices:
-                    continue
-                delta = getattr(choices[0].delta, "content", None)
-                if delta:
-                    yield delta
+            async with stream:
+                async for chunk in stream:
+                    choices = getattr(chunk, "choices", None) or []
+                    if not choices:
+                        continue
+                    delta = getattr(choices[0].delta, "content", None)
+                    if delta:
+                        yield delta
             return
         except Exception as exc:
             last_error = exc

api/bigrag/services/chat/readiness.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from bigrag.models.chat_readiness import ChatReadinessResponse
+from bigrag.services.chat.turn.credentials import (
+    ChatCredentialDecryptionError,
+    ChatCredentialDestinationError,
+    assert_credentials_allowed_for_base_url,
+    resolve_stored_chat_credentials,
+)
+from bigrag.services.runtime_settings import get_values
+from bigrag.services.url_security import UnsafeOutboundUrlError, normalize_url_root
+
+
+async def get_chat_readiness(session: AsyncSession, user: dict) -> ChatReadinessResponse:
+    runtime = await get_values(["chat_base_url"])
+    try:
+        credentials = await resolve_stored_chat_credentials(session, user)
+    except ChatCredentialDecryptionError:
+        return ChatReadinessResponse(credential_ready=False, reason="decryption_unavailable")
+    if not credentials:
+        return ChatReadinessResponse(credential_ready=False, reason="missing_credentials")
+    try:
+        raw_base_url = runtime["chat_base_url"]
+        base_url = normalize_url_root(raw_base_url) if raw_base_url else None
+    except (UnsafeOutboundUrlError, ValueError):
+        return ChatReadinessResponse(credential_ready=False, reason="credential_blocked")
+    try:
+        assert_credentials_allowed_for_base_url(credentials, base_url, request_base_url=None)
+    except ChatCredentialDestinationError:
+        return ChatReadinessResponse(credential_ready=False, reason="credential_blocked")
+    return ChatReadinessResponse(
+        credential_ready=True,
+        credential_source="saved" if credentials[0].source == "saved chat key" else "instance",
+    )