fix(serve): cap on-demand renders per /search request

serve/src/pixelrag_serve/api.py

@@ -279,6 +279,10 @@ class StatusResponse(BaseModel):
 _MAX_IMAGE_B64_LEN = 10 * 1024 * 1024
 _MAX_IMAGE_PIXELS = 25_000_000
 
+# Most on-demand renders one /search request may trigger (see the search handler).
+# A real caller asks for n_docs=10 with include_images, so this is generous.
+_MAX_ONDEMAND_RENDERS = 16
+
 
 def _parse_queries(
     queries: list[Query], instruction: str | None = None
@@ -540,6 +544,18 @@ async def search(req: SearchRequest, request: Request):
     # Build results
     tiles_dir = _state.get("tiles_dir", "")
 
+    # Budget for on-demand renders across the whole request. A render is a
+    # Chrome page load under a process-global lock, bounded only by
+    # PIXELRAG_RENDER_TIMEOUT (120s by default), and nothing else caps how many
+    # one request can ask for: at this endpoint's own limits (32 queries,
+    # n_docs <= 1000) a single include_images request could queue tens of
+    # thousands of them, stalling every other caller's renders behind the same
+    # lock. Past the budget a hit is returned without its image — what a failed
+    # render or a missing tile already yields — and the client can still fetch
+    # it from /tile. Cache hits are cheap and don't spend the budget.
+    ondemand = _state.get("ondemand")
+    renders_left = _MAX_ONDEMAND_RENDERS
+
     results = []
     for qi in range(len(req.queries)):
         hits = []
@@ -559,11 +575,17 @@ async def search(req: SearchRequest, request: Request):
             if req.include_images and tile_path and os.path.exists(tile_path):
                 with open(tile_path, "rb") as fp:
                     img_b64 = base64.b64encode(fp.read()).decode()
-            elif req.include_images and _state.get("ondemand") is not None:
-                # Render off the event loop: _ondemand_chunk_b64 -> render_url uses
-                # asyncio.run(), which raises "cannot be called from a running event
-                # loop" if invoked directly here. Offload to a worker thread.
-                img_b64 = await asyncio.to_thread(_ondemand_chunk_b64, aid, ti, ci, th)
+            elif req.include_images and ondemand is not None:
+                cached = ondemand.cached_chunk_path(aid, ti, ci) is not None
+                if cached or renders_left > 0:
+                    if not cached:
+                        renders_left -= 1
+                    # Render off the event loop: _ondemand_chunk_b64 -> render_url uses
+                    # asyncio.run(), which raises "cannot be called from a running event
+                    # loop" if invoked directly here. Offload to a worker thread.
+                    img_b64 = await asyncio.to_thread(
+                        _ondemand_chunk_b64, aid, ti, ci, th
+                    )
             # Expose a relative tile path, not the absolute server filesystem
             # path (avoids leaking the host's directory layout; clients fetch
             # tiles via /tile/{article_id}/{tile_index}/{chunk_index}).

serve/src/pixelrag_serve/render_ondemand.py

@@ -112,14 +112,27 @@ def _ensure_chrome(self) -> str:
     def _article_dir(self, article_id: int) -> str:
         return os.path.join(self.cache_dir, f"{article_id}.png.tiles")
 
+    def _chunk_file(self, article_id: int, tile_index: int, chunk_index: int) -> str:
+        chunk_name = f"chunk_{tile_index:04d}_{chunk_index:02d}.png"
+        return os.path.join(self._article_dir(article_id), chunk_name)
+
+    def cached_chunk_path(self, article_id: int, tile_index: int, chunk_index: int):
+        """Path to an already-rendered chunk, or None. Never renders.
+
+        Lets a caller tell a cheap cache hit from a page render before
+        committing to one.
+        """
+        cpath = self._chunk_file(article_id, tile_index, chunk_index)
+        return cpath if os.path.exists(cpath) else None
+
     def chunk_path(
         self, article_id: int, title: str, tile_index: int, chunk_index: int
     ):
         """Path to chunk_{ti}_{ci}.png, rendering+chunking the page on a cache miss."""
-        chunk_name = f"chunk_{tile_index:04d}_{chunk_index:02d}.png"
-        cpath = os.path.join(self._article_dir(article_id), chunk_name)
-        if os.path.exists(cpath):
-            return cpath
+        cached = self.cached_chunk_path(article_id, tile_index, chunk_index)
+        if cached:
+            return cached
+        cpath = self._chunk_file(article_id, tile_index, chunk_index)
         if not title:
             return None
         with _render_lock:

tests/test_serve_ondemand_budget.py

@@ -0,0 +1,124 @@
+"""Cap on on-demand renders per /search request (DoS hardening).
+
+An on-demand render is a Chrome page load under a process-global lock with a
+120s timeout. Without a per-request budget, one `include_images` request could
+queue one per hit — at the endpoint's own limits, tens of thousands — and every
+other caller's renders stall behind the same lock.
+
+Driven through the real handler with pre-computed embeddings, so no model or
+index is needed.
+"""
+
+import numpy as np
+import pytest
+
+pytest.importorskip("fastapi", reason="serve extra not installed")
+from fastapi.testclient import TestClient
+from pixelrag_serve import api
+
+DIM = 4
+N_HITS = 40  # more hits than the budget allows renders
+
+
+class _CountingOnDemand:
+    """Stands in for OnDemandTiles: counts renders, never touches Chrome."""
+
+    def __init__(self, cached: set[int] | None = None):
+        self.cached = cached or set()
+        self.renders = 0
+
+    def cached_chunk_path(self, article_id, tile_index, chunk_index):
+        return f"/cached/{article_id}.png" if article_id in self.cached else None
+
+    def chunk_path(self, article_id, title, tile_index, chunk_index):
+        if article_id not in self.cached:
+            self.renders += 1
+
+
+class _StubBackend:
+    dimension = DIM
+
+    def set_nprobe(self, n):
+        pass
+
+    def reset_nprobe(self):
+        pass
+
+    def raw_search(self, query_vectors, k, **kw):
+        return [
+            [
+                {
+                    "score": 1.0,
+                    "vector_id": i,
+                    "article_id": i,
+                    "tile_index": 0,
+                    "chunk_index": 0,
+                    "y_offset": 0,
+                    "tile_height": 100,
+                }
+                for i in range(min(k, N_HITS))
+            ]
+            for _ in query_vectors
+        ]
+
+
+@pytest.fixture
+def client(tmp_path):
+    def _make(ondemand):
+        api._state.clear()
+        api._state.update(
+            {
+                "backend": _StubBackend(),
+                "articles": [f"Article_{i}" for i in range(N_HITS)],
+                "dimension": DIM,
+                "tiles_dir": str(tmp_path),  # empty: nothing is on disk
+                "ondemand": ondemand,
+            }
+        )
+        api._article_pages.cache_clear()
+        return TestClient(api.app)
+
+    yield _make
+    api._state.clear()
+
+
+def _search(client, n_docs):
+    body = {
+        "queries": [{"embedding": np.ones(DIM).tolist()}],
+        "n_docs": n_docs,
+        "include_images": True,
+    }
+    resp = client.post("/search", json=body)
+    assert resp.status_code == 200, resp.text
+    return resp.json()
+
+
+def test_renders_are_capped_per_request(client):
+    od = _CountingOnDemand()
+    body = _search(client(od), N_HITS)
+    # Every hit still comes back — only the images past the budget are dropped.
+    assert len(body["results"][0]["hits"]) == N_HITS
+    assert od.renders == api._MAX_ONDEMAND_RENDERS
+
+
+def test_budget_spans_all_queries_in_one_request(client):
+    od = _CountingOnDemand()
+    c = client(od)
+    resp = c.post(
+        "/search",
+        json={
+            "queries": [{"embedding": np.ones(DIM).tolist()} for _ in range(4)],
+            "n_docs": N_HITS,
+            "include_images": True,
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    # One budget for the request, not one per query.
+    assert od.renders == api._MAX_ONDEMAND_RENDERS
+
+
+def test_cache_hits_do_not_spend_the_budget(client):
+    # Every article already rendered: serving them must cost no renders at all.
+    od = _CountingOnDemand(cached=set(range(N_HITS)))
+    _search(client(od), N_HITS)
+    assert od.renders == 0