From 136c984a03ed491a3139717c915bd922bfdd0c90 Mon Sep 17 00:00:00 2001
From: Vito Sansevero <vito.sansevero@gmail.com>
Date: Tue, 27 Oct 2026 12:39:58 -0700
Subject: [PATCH] test(search): pin prompt-only tags and multi-chunk tag
 lookups (#379)

Reads the tag names and the locators in the same annotation read helper.
---
 database/annotations.py            | 12 ++++++------
 tests/test_annotations.py          | 14 ++++++++++++--
 tests/test_api_workspace_images.py | 21 +++++++++++++++++++++
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/database/annotations.py b/database/annotations.py
index d7b4b681..96062fa6 100644
--- a/database/annotations.py
+++ b/database/annotations.py
@@ -454,12 +454,12 @@ class AnnotationOperations:
         already folded). One read of the names that tag an image.
         """
         wanted = set(wanted)
-        conn = self.model.get_connection()
-        return [
-            tag_id
-            for tag_id, name in conn.execute(_IMAGE_TAG_NAMES)
-            if name in wanted or (query and query in name.casefold())
-        ]
+        with self._annotation_read() as conn:
+            return [
+                tag_id
+                for tag_id, name in conn.execute(_IMAGE_TAG_NAMES)
+                if name in wanted or (query and query in name.casefold())
+            ]
 
     def image_tag_locators(self, tag_ids) -> set[str]:
         """Locators with any byte record (current or not) carrying one of
diff --git a/tests/test_annotations.py b/tests/test_annotations.py
index 0d3257cf..87ac4d5c 100644
--- a/tests/test_annotations.py
+++ b/tests/test_annotations.py
@@ -553,8 +553,8 @@ class CurrentImageTagsTests(AnnotationTestCase):
         self.assertEqual(current, {locators[0]: ["z", "y"], locators[502]: ["x"]})
         lookups = [s for s in statements if "media_annotation_files" in s]
         self.assertEqual(len(lookups), 2)
-        # The tag search finds every record's file by tag id, the empty set's
-        # file never, and each set keeps its position order.
+        # The tag search finds every record's file by tag id, and the empty
+        # set's file never.
         self.assertEqual(
             self.db.image_tag_locators(self.db.image_tag_ids(["z", "x"], "")),
             {locators[0], locators[502]},
@@ -564,6 +564,16 @@ class CurrentImageTagsTests(AnnotationTestCase):
             {locators[250]: []},
         )
 
+    def test_tag_locators_cover_more_ids_than_one_chunk(self):
+        # A broad text search can match more than 500 names: every id's file
+        # comes back across the chunks.
+        locators = [self.locator(f"t{i:04}.png") for i in range(503)]
+        for i, loc in enumerate(locators):
+            self.replace(loc, tags=[f"name_{i:04}"])
+        ids = self.db.image_tag_ids([], "name_")
+        self.assertEqual(len(ids), 503)
+        self.assertEqual(self.db.image_tag_locators(ids), set(locators))
+
 
 class AnnotationIsolationTests(AnnotationTestCase):
     def test_cascade_from_vocabulary_delete_keeps_byte_records(self):
diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py
index 9fb3b7cf..cb2ff6dd 100644
--- a/tests/test_api_workspace_images.py
+++ b/tests/test_api_workspace_images.py
@@ -2163,6 +2163,27 @@ class TestWorkspaceImageTags(ImageAPITestCase):
                 "tag_id IN" in sql or "EXISTS" in sql or "byte_id IN" in sql, sql
             )
 
+    async def test_the_image_tag_lookup_never_reads_prompt_tags(self):
+        self.seed_differential()
+        linked = self.media("linked.png", mtime=105)
+        prompt = self._save_prompt(
+            "Prompt with its own tag", tags=["prompt_only_beach"]
+        )
+        self.api.db.link_image_to_prompt(prompt, str(linked))
+        db = self.api.db
+        # The tag is only on a prompt: the image-tag lookup never yields it,
+        # under the Tags filter or the text search.
+        self.assertEqual(db.image_tag_ids(["prompt_only_beach"], ""), [])
+        self.assertEqual(db.image_tag_ids([], "prompt_only"), [])
+        self.assertEqual(
+            db.image_tag_locators(db.image_tag_ids([], "beach")),
+            reference_candidates(db, [], "beach"),
+        )
+        # The search still matches through the linked positive's prompt tags
+        # (the design keeps that, #379 "Behavior that does not change").
+        _, body = await self.get(q="prompt_only")
+        self.assertEqual(self.names(body), ["linked.png"])
+
     @unittest.skipUnless(os.environ.get("PM_PERF") == "1", "timing run (PM_PERF=1)")
     async def test_tag_search_timing(self):
         """Linux timing for the PR body: 23,600 records, 10 tags each."""