fix(library): preserve full name order and batch previews (#208)

database/library.py

@@ -6,7 +6,7 @@ _ELIGIBLE = """(EXISTS(SELECT 1 FROM prompt_pairs ep WHERE ep.positive_prompt_id
 _ORDERS = {
     "newest": ("COALESCE(julianday(p.created_at),0)", "DESC"),
     "oldest": ("COALESCE(julianday(p.created_at),0)", "ASC"),
-    "name": ("LOWER(SUBSTR(p.text,1,160))", "ASC"),
+    "name": ("LOWER(p.text)", "ASC"),
     "rating": ("COALESCE(p.rating,0)", "DESC"),
 }
 
@@ -249,6 +249,30 @@ class LibraryOperations:
         )
         return {"categories": [row["category"] for row in rows]}
 
+    def get_library_previews(self, positive_ids: list[int]) -> dict[int, list[dict]]:
+        """Read at most two links per parent in one bounded page query."""
+        if (
+            not isinstance(positive_ids, list)
+            or len(positive_ids) > 200
+            or any(type(value) is not int or value <= 0 for value in positive_ids)
+        ):
+            raise ValueError("Invalid preview parent IDs")
+        if not positive_ids:
+            return {}
+        result = {value: [] for value in positive_ids}
+        marks = ",".join("?" for _ in positive_ids)
+        rows = self.model.get_connection().execute(
+            "WITH ranked AS (SELECT id,prompt_id,image_path,ROW_NUMBER() OVER "
+            "(PARTITION BY prompt_id ORDER BY id DESC) AS position "
+            f"FROM generated_images WHERE prompt_id IN ({marks})) "
+            "SELECT id,prompt_id,image_path FROM ranked WHERE position<=2 "
+            "ORDER BY prompt_id,id DESC",
+            positive_ids,
+        )
+        for row in rows:
+            result[row["prompt_id"]].append(dict(row))
+        return result
+
     def list_library_images(
         self,
         positive_id: int,

py/api/library.py

@@ -6,6 +6,9 @@ import hashlib
 import hmac
 import json
 import os
+import time
+from collections import OrderedDict
+from collections.abc import Callable
 from pathlib import Path
 
 from aiohttp import web
@@ -16,16 +19,73 @@ from .pairs import pair_response
 _SECRET = os.urandom(32)
 
 
+class _NameCursorStore:
+    """Full lexical keys stay off URLs, bounded by age, count and UTF-8 bytes."""
+
+    def __init__(
+        self,
+        ttl: int = 3600,
+        max_entries: int = 128,
+        max_bytes: int = 8 * 1024 * 1024,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self.clock = clock
+        self.ttl = ttl
+        self.max_entries = max_entries
+        self.max_bytes = max_bytes
+        self.entries = OrderedDict()
+        self.size = 0
+
+    def _remove(self, reference: str) -> None:
+        _expires, _scope, raw = self.entries.pop(reference)
+        self.size -= len(raw)
+
+    def _prune(self) -> None:
+        now = self.clock()
+        for reference, (expires, _scope, _raw) in list(self.entries.items()):
+            if expires <= now:
+                self._remove(reference)
+
+    def put(self, key: object, scope: str) -> str:
+        raw = json.dumps(key, ensure_ascii=False, allow_nan=False).encode()
+        if len(raw) > min(512 * 1024, self.max_bytes):
+            raise ValueError("Prompt name exceeds the cursor limit")
+        self._prune()
+        while self.entries and (
+            len(self.entries) >= self.max_entries
+            or self.size + len(raw) > self.max_bytes
+        ):
+            self._remove(next(iter(self.entries)))
+        reference = base64.urlsafe_b64encode(os.urandom(24)).decode()
+        self.entries[reference] = (self.clock() + self.ttl, scope, raw)
+        self.size += len(raw)
+        return reference
+
+    def get(self, reference: str, scope: str) -> object:
+        self._prune()
+        entry = self.entries.get(reference)
+        if entry is None or entry[1] != scope:
+            raise ValueError("Name cursor expired; refresh the library")
+        self.entries.move_to_end(reference)
+        return json.loads(entry[2])
+
+
+_NAME_CURSORS = _NameCursorStore()
+
+
 def _scope(value: dict) -> str:
     return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest()
 
 
-def _encode(key: object, scope: str) -> str | None:
+def _encode(key: object, scope: str, *, name: bool = False) -> str | None:
     if key is None:
         return None
-    payload = json.dumps(
-        {"scope": scope, "key": key}, separators=(",", ":"), allow_nan=False
-    ).encode()
+    value = (
+        {"scope": scope, "reference": _NAME_CURSORS.put(key, scope)}
+        if name
+        else {"scope": scope, "key": key}
+    )
+    payload = json.dumps(value, separators=(",", ":"), allow_nan=False).encode()
     signature = hmac.new(_SECRET, payload, hashlib.sha256).digest()
     return base64.urlsafe_b64encode(signature + payload).decode().rstrip("=")
 
@@ -45,8 +105,13 @@ def _decode(token: str | None, scope: str) -> object:
         ):
             raise ValueError
         value = json.loads(payload)
-        if set(value) != {"scope", "key"} or value["scope"] != scope:
+        if (
+            set(value) not in ({"scope", "key"}, {"scope", "reference"})
+            or value["scope"] != scope
+        ):
             raise ValueError
+        if "reference" in value:
+            return _NAME_CURSORS.get(value["reference"], scope)
         return value["key"]
     except (ValueError, TypeError, binascii.Error, UnicodeError):
         raise ValueError("Invalid library cursor; refresh the library") from None
@@ -124,12 +189,16 @@ class LibraryRoutesMixin:
         roots = list(self._get_all_output_dirs())
 
         def previews() -> None:
+            records = self.db.get_library_previews(
+                [parent["id"] for parent in result["items"]]
+            )
             for parent in result["items"]:
-                images = self.db.list_library_images(parent["id"], limit=2)
-                parent["previews"] = self._library_media(images["images"], roots)
+                parent["previews"] = self._library_media(records[parent["id"]], roots)
 
         await self._run_in_executor(previews)
-        result["next_cursor"] = _encode(result.pop("next_key"), scope)
+        result["next_cursor"] = _encode(
+            result.pop("next_key"), scope, name=options["sort"] == "name"
+        )
         return web.json_response({"success": True, **result})
 
     @pair_response

tests/test_api_library.py

@@ -1,6 +1,7 @@
 """Native library endpoints use bounded cursors and rooted media identities."""
 
 import json
+import unittest
 from unittest.mock import patch
 
 try:
@@ -178,3 +179,106 @@ class TestLibraryAPI(ImageAPITestCase):
         self.assertEqual(
             [item["text"] for item in literal["items"]], ["literal % and _"]
         )
+
+    async def test_maximum_parent_page_uses_bounded_sql_queries_for_previews(self):
+        for index in range(200):
+            self._save_prompt(f"Parent {index}")
+        statements = []
+        original = self.api.db.model.get_connection
+
+        def trace_connection():
+            connection = original()
+            connection.set_trace_callback(statements.append)
+            return connection
+
+        with patch.object(
+            self.api.db.model, "get_connection", side_effect=trace_connection
+        ):
+            response, result = await self.get(limit=200)
+        self.assertEqual(response.status, 200)
+        self.assertEqual(len(result["items"]), 200)
+        selects = [
+            query
+            for query in statements
+            if query.lstrip().upper().startswith(("SELECT", "WITH"))
+        ]
+        self.assertLessEqual(len(selects), 5)
+
+    async def test_long_name_cursor_stays_small_and_preserves_full_sorting(self):
+        prefix = "猫" * 10_000
+        self._save_prompt(prefix + "Z")
+        alpha = self._save_prompt(prefix + "A")
+        _, first = await self.get(sort="name", limit=1)
+        self.assertEqual(first["items"][0]["id"], alpha)
+        self.assertLess(len(first["next_cursor"]), 512)
+        with self.api.db.model.get_connection() as connection:
+            connection.execute("DELETE FROM prompts WHERE id=?", (alpha,))
+        _, second = await self.get(sort="name", cursor=first["next_cursor"])
+        self.assertEqual([item["text"] for item in second["items"]], [prefix + "Z"])
+
+    async def test_name_cursor_expiry_and_eviction_require_refresh_without_large_urls(
+        self,
+    ):
+        from py.api import library as module
+
+        for text in ("Alpha", "Beta", "Gamma"):
+            self._save_prompt(text)
+        now = [100]
+        store = module._NameCursorStore(max_entries=1, clock=lambda: now[0])
+        with patch.object(module, "_NAME_CURSORS", store):
+            _, first = await self.get(sort="name", limit=1)
+            now[0] = 3699
+            response, _ = await self.get(sort="name", cursor=first["next_cursor"])
+            self.assertEqual(response.status, 200)
+            now[0] = 3700
+            response, result = await self.get(sort="name", cursor=first["next_cursor"])
+            self.assertEqual(response.status, 400)
+            self.assertIn("refresh", result["error"])
+            response, fresh = await self.get(sort="name", limit=1)
+            self.assertEqual(response.status, 200)
+            _, replacement = await self.get(sort="name", limit=1)
+            self.assertNotEqual(fresh["next_cursor"], replacement["next_cursor"])
+            response, result = await self.get(sort="name", cursor=fresh["next_cursor"])
+            self.assertEqual(response.status, 400)
+            self.assertIn("refresh", result["error"])
+            response, _ = await self.get(
+                sort="name", cursor=replacement["next_cursor"], text="Alpha"
+            )
+            self.assertEqual(response.status, 400)
+
+    async def test_batch_preview_hydration_keeps_two_real_images_per_parent(self):
+        for parent_index in range(2):
+            positive = self._save_prompt(f"Parent {parent_index}")
+            for index in range(3):
+                path = make_png(self.output_dir / f"{parent_index}-{index}.png")
+                self.api.db.link_image_to_prompt(positive, str(path))
+        _, result = await self.get()
+        self.assertEqual(len(result["items"]), 2)
+        for item in result["items"]:
+            self.assertEqual(len(item["previews"]), 2)
+            self.assertTrue(item["previews"][0]["path"].endswith("-2.png"))
+            self.assertTrue(item["previews"][1]["path"].endswith("-1.png"))
+            self.assertTrue(all(image["root_index"] == 0 for image in item["previews"]))
+
+
+class TestNameCursorBudget(unittest.TestCase):
+    def test_lru_byte_budget_scope_and_oversize_keys_are_bounded(self):
+        from py.api.library import _NameCursorStore
+
+        store = _NameCursorStore(max_entries=2, max_bytes=40)
+        one = store.put("one", "scope")
+        two = store.put("two", "scope")
+        self.assertEqual(store.get(one, "scope"), "one")
+        three = store.put("three", "scope")
+        with self.assertRaises(ValueError):
+            store.get(two, "scope")
+        self.assertEqual(store.get(three, "scope"), "three")
+        with self.assertRaises(ValueError):
+            store.get(one, "other")
+        self.assertLessEqual(store.size, store.max_bytes)
+        larger = store.put("x" * 35, "scope")
+        self.assertEqual(len(store.entries), 1)
+        self.assertEqual(store.get(larger, "scope"), "x" * 35)
+        with self.assertRaises(ValueError):
+            store.put("x" * 41, "scope")
+        self.assertEqual(len(store.entries), 1)

tests/test_library.py

@@ -173,6 +173,7 @@ class TestLibrary(unittest.TestCase):
             self.db.list_library()
             self.db.get_library_positive(b)
             self.db.get_library_options()
+            self.db.get_library_previews([b])
             self.db.list_library_images(b)
         finally:
             connection.set_authorizer(None)
@@ -244,3 +245,55 @@ class TestLibrary(unittest.TestCase):
         self.assertTrue(
             all(self.db.list_library(category=value)["total"] == 1 for value in options)
         )
+
+    def test_name_sort_uses_full_long_text_and_keeps_deleted_anchor_key(self):
+        prefix = "shared prefix " * 30
+        self.save(prefix + "Zulu")
+        alpha = self.save(prefix + "Alpha")
+        self.save(prefix + "Middle")
+        first = self.db.list_library(sort="name", limit=1)
+        self.assertEqual(first["items"][0]["id"], alpha)
+        self.assertEqual(first["next_key"]["value"], (prefix + "Alpha").lower())
+        with self.db.model.get_connection() as connection:
+            connection.execute("DELETE FROM prompts WHERE id=?", (alpha,))
+        rest = self.db.list_library(sort="name", after=first["next_key"])
+        self.assertEqual(
+            [item["text"] for item in rest["items"]],
+            [prefix + "Middle", prefix + "Zulu"],
+        )
+
+    def test_batch_previews_return_two_newest_links_per_parent_in_one_select(self):
+        parents = [self.save(f"Parent {index}") for index in range(200)]
+        expected = {}
+        for parent in parents[:3]:
+            ids = [
+                self.db.link_image_to_prompt(
+                    parent, str(Path(self.temp.name) / f"{parent}-{index}.png")
+                )
+                for index in range(3)
+            ]
+            expected[parent] = ids[-1:0:-1]
+        connection = self.db.model.get_connection()
+        statements = []
+        connection.set_trace_callback(statements.append)
+        try:
+            previews = self.db.get_library_previews(parents)
+        finally:
+            connection.set_trace_callback(None)
+        selects = [
+            query
+            for query in statements
+            if query.lstrip().upper().startswith(("SELECT", "WITH"))
+        ]
+        self.assertEqual(len(selects), 1)
+        self.assertEqual(
+            {
+                parent: [image["id"] for image in images]
+                for parent, images in previews.items()
+                if images
+            },
+            expected,
+        )
+        self.assertEqual(self.db.get_library_previews([]), {})
+        with self.assertRaises(ValueError):
+            self.db.get_library_previews(parents + [9999])