unsloth/studio/backend/core/rag/authorization.py
Daniel Han ab0828b976 Studio: fix RAG correctness bugs
Backend:
- Deterministic SQLite connection cleanup. The RAG code used bare
  `with get_connection() as conn:`, which commits but never closes, leaning
  on GC to release handles (the rest of studio_db closes explicitly). Add a
  closing_connection() context manager that commits/rolls back like sqlite3's
  own manager and always closes, and route all 30 RAG call sites through it.
- filter_by_min_score no longer drops BM25-only and figure-ref hits. min_score
  is a cosine floor, so it now gates only hits that carry a dense_score;
  lexical and figure-ref hits (dense_score is None) pass through instead of
  being silently discarded when the floor is raised.
- Fix two tests that could not pass against the production code: the RRF
  fusion test asserted the wrong winner (c edges out b: 0.032266 vs 0.032258),
  and two tool-handler scope tests stubbed retrieve_hybrid without accepting
  the embedder_model kwarg the handler now passes (TypeError was swallowed,
  leaving captured["scope"] unset).

Frontend:
- Removing an in-flight upload chip now routes through the teardown thunk
  already registered for the aggregate-progress toast (abort, unsubscribe,
  release the index slot, delete the backend doc with the correct kb/thread
  scope key it closed over) and clears the toast entry. Deleting directly
  leaked the concurrency slot and hardcoded the thread scope, mis-targeting
  KB-scoped docs. Applied in both the composer hook and the compare-view
  composer; drop the now-vestigial chip-scope-key tracking and unused
  activeThreadId selectors. Add index-progress-store.remove(id).
2026-05-31 09:56:23 +00:00

97 lines
4 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Subject-scoped authorization for RAG document preview routes.
Used by `/api/rag/documents/{document_id}/file` and
`/api/rag/documents/{document_id}/preview-target` to enforce that the
current authenticated subject is allowed to see a given document and
chunk. Existence and authorization failures collapse to a single 404
so the API does not leak document IDs to a non-owner.
"""
from __future__ import annotations
import sqlite3
from fastapi import HTTPException
from storage.studio_db import closing_connection
_NOT_FOUND_DETAIL = "Document not found"
def document_for_subject_or_404(
document_id: str,
current_subject: str,
) -> sqlite3.Row:
"""Return the `rag_documents` row if `current_subject` may access it.
Authorization rules:
- KB documents: the document's KB must have
`rag_knowledge_bases.owner_user_id == current_subject`. A KB with a
NULL owner is not accessible through this helper (legacy pre-auth
rows must be migrated or accessed via admin tooling).
- Thread documents: thread-scoped RAG documents are gated by an
explicit single-user invariant for Studio's current release. The
`chat_threads` table does not yet carry an `owner_user_id` column,
so we cannot bind a thread to a specific subject in the schema.
The helper still requires (a) an authenticated subject (enforced
by the route's `Depends(get_current_subject)`) and (b) that the
referenced thread actually exists in `chat_threads`. A missing
thread row collapses to 404 so a non-existent thread cannot
silently grant access through a dangling `thread_id`.
# TODO(thread-owner): once `chat_threads.owner_user_id` exists, join
# through it like KB docs and drop the single-user invariant. Update
# `tests/test_rag_authorization.py::test_thread_doc_other_user_404`
# to assert per-user isolation rather than thread existence.
Both not-found and not-authorized raise `HTTPException(404)` with the
same detail string. Callers must NOT distinguish the two cases in
their response, to avoid leaking document existence to a non-owner.
Returns the document row so the caller can read `stored_path`,
`filename`, `content_type`, etc. without re-querying.
"""
if not document_id or not current_subject:
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
with closing_connection() as conn:
row = conn.execute(
"SELECT * FROM rag_documents WHERE id = ?",
(document_id,),
).fetchone()
if row is None:
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
kb_id = row["kb_id"]
thread_id = row["thread_id"]
if kb_id is not None:
owner_row = conn.execute(
"SELECT owner_user_id FROM rag_knowledge_bases WHERE id = ?",
(kb_id,),
).fetchone()
if owner_row is None:
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
owner = owner_row["owner_user_id"]
if owner is None or owner != current_subject:
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
return row
if thread_id is not None:
# Single-user invariant (see TODO above): require the thread row to
# exist; an unknown thread_id is not-found, not a silent grant.
thread_row = conn.execute(
"SELECT id FROM chat_threads WHERE id = ?",
(thread_id,),
).fetchone()
if thread_row is None:
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
return row
# Docs must belong to a KB or a thread (DB CHECK enforces XOR on insert);
# a row satisfying neither is corrupt — treat as 404.
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)