From 16f1427b7ac3095f9456db2162e79cc55ddad7d3 Mon Sep 17 00:00:00 2001
From: stumpylog <797416+stumpylog@users.noreply.github.com>
Date: Thu, 27 Aug 2026 11:37:07 -0700
Subject: [PATCH] Not bad catches from Copilot, if a little extra secure

---
 src/paperless_ai/ai_classifier.py            | 30 +++++++++++++++-----
 src/paperless_ai/tests/test_ai_classifier.py |  3 +-
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py
index 7b7c3a85d..d2a0ab32f 100644
--- a/src/paperless_ai/ai_classifier.py
+++ b/src/paperless_ai/ai_classifier.py
@@ -6,6 +6,8 @@ from django.contrib.auth.models import User
 
 from documents.models import Document
 from documents.permissions import permitted_object_ids
+from documents.permissions import restrict_queryset_to_visible
+from documents.permissions import user_is_unrestricted
 from paperless.config import AIConfig
 from paperless_ai.base_model import ClassificationSuggestions
 from paperless_ai.base_model import TaxonomyChoiceDict
@@ -50,19 +52,32 @@ def _fulltext_similar_documents(
     similarity - cruder, but far better than no candidates at all.
     more_like_this_ids returns only a ranked ID list, no scores, so weight is
     synthesized from rank (descending from top_k) rather than claiming a
-    similarity magnitude that doesn't exist. A superuser is normalized to
-    ``None`` before calling, since the backend's permission filter has no
-    superuser short-circuit of its own.
+    similarity magnitude that doesn't exist. An unrestricted user (none, or an
+    active superuser - see user_is_unrestricted) is normalized to ``None``
+    before calling, since the backend's permission filter has no superuser
+    short-circuit of its own. Results are re-checked with
+    restrict_queryset_to_visible() since Tantivy's indexed permission fields
+    lag the DB via async reindexing.
     """
     from documents.search import get_backend
 
-    search_user = None if user is not None and user.is_superuser else user
+    unrestricted = user_is_unrestricted(user)
+    search_user = None if unrestricted else user
     backend = get_backend()
     similar_ids = backend.more_like_this_ids(
         document.pk,
         user=search_user,
         limit=top_k,
     )
+    if not unrestricted:
+        allowed_ids = set(
+            restrict_queryset_to_visible(
+                Document.objects.filter(pk__in=similar_ids),
+                user,
+                "view_document",
+            ).values_list("pk", flat=True),
+        )
+        similar_ids = [doc_id for doc_id in similar_ids if doc_id in allowed_ids]
     return [
         SimilarDocument(document_id=doc_id, weight=float(top_k - rank))
         for rank, doc_id in enumerate(similar_ids)
@@ -221,9 +236,10 @@ def get_taxonomy_context(
         candidates = build_taxonomy_candidates(similar_documents, user)
 
         similar_doc_ids = [s["document_id"] for s in similar_documents]
-        similar_docs = list(
-            Document.objects.filter(pk__in=similar_doc_ids)[:max_docs],
-        )
+        docs_by_id = Document.objects.in_bulk(similar_doc_ids)
+        similar_docs = [
+            docs_by_id[doc_id] for doc_id in similar_doc_ids if doc_id in docs_by_id
+        ][:max_docs]
         context_blocks = []
         for similar in similar_docs:
             text = similar.content[:1000] or ""
diff --git a/src/paperless_ai/tests/test_ai_classifier.py b/src/paperless_ai/tests/test_ai_classifier.py
index 289beb8c5..2261dc9e1 100644
--- a/src/paperless_ai/tests/test_ai_classifier.py
+++ b/src/paperless_ai/tests/test_ai_classifier.py
@@ -634,6 +634,7 @@ class TestFulltextSimilarDocuments:
             - At most 2 results are returned
         """
         source = DocumentFactory.create(content="shared overlapping keyword text")
+        fulltext_backend.add_or_update(source)
         for _ in range(4):
             fulltext_backend.add_or_update(
                 DocumentFactory.create(content="shared overlapping keyword text"),
@@ -641,7 +642,7 @@ class TestFulltextSimilarDocuments:
 
         result = _fulltext_similar_documents(source, user=None, top_k=2)
 
-        assert len(result) <= 2
+        assert len(result) == 2
 
     def test_result_shape_is_similar_document(
         self,