feat: cap SSE stream at PAGE_SIZE and buffer remaining results

Co-Authored-By: Claude Opus 4.6 <[email protected]>

tests/test_app.py

@@ -1,10 +1,12 @@
 import io
+import json
 from pathlib import Path
 
 import pytest
 from httpx import ASGITransport, AsyncClient
 
-from unredact.app import app
+import unredact.app
+from unredact.app import app, PAGE_SIZE, _solve_results
 
 
 @pytest.fixture
@@ -261,3 +263,125 @@ async def test_spot_returns_analysis(pdf_bytes: bytes):
             assert "segments" in data["analysis"]
             assert "gap" in data["analysis"]
             assert "line" in data["analysis"]
+
+
+@pytest.mark.anyio
+async def test_solve_paginates_results():
+    """Solve should stream at most PAGE_SIZE results and buffer the rest."""
+    from unredact.app import PAGE_SIZE
+
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://test") as client:
+        resp = await client.get("/api/fonts")
+        fonts = resp.json()["fonts"]
+        font = next(f for f in fonts if f["available"])
+
+        # Use word mode with full vocab and wide tolerance to get many results
+        resp = await client.post("/api/solve", json={
+            "font_id": font["id"],
+            "font_size": 40,
+            "gap_width_px": 80.0,
+            "tolerance_px": 10.0,
+            "left_context": "",
+            "right_context": "",
+            "hints": {"charset": "lowercase"},
+            "mode": "word",
+            "vocab_size": 0,
+        })
+        assert resp.status_code == 200
+
+        # Parse SSE events
+        events = []
+        for line in resp.text.split("\n"):
+            if line.startswith("data: "):
+                try:
+                    events.append(json.loads(line[6:]))
+                except (json.JSONDecodeError, ValueError):
+                    pass
+
+        matches = [e for e in events if e.get("status") == "match"]
+        done_events = [e for e in events if e.get("status") == "done"]
+        page_events = [e for e in events if e.get("status") == "page_complete"]
+
+        assert len(done_events) == 1
+        total = done_events[0]["total_found"]
+
+        if total > PAGE_SIZE:
+            assert len(matches) == PAGE_SIZE
+            assert len(page_events) == 1
+            assert page_events[0]["sent"] == PAGE_SIZE
+        else:
+            assert len(matches) == total
+            assert len(page_events) == 0
+
+
+def _parse_sse_events(body: str) -> list[dict]:
+    """Parse SSE body text into a list of JSON event dicts."""
+    events = []
+    for line in body.splitlines():
+        line = line.strip()
+        if line.startswith("data: "):
+            try:
+                events.append(json.loads(line[6:]))
+            except (json.JSONDecodeError, ValueError):
+                continue
+    return events
+
+
+@pytest.mark.anyio
+async def test_solve_pagination_caps_at_page_size(monkeypatch):
+    """POST /api/solve should cap streamed match events at PAGE_SIZE and send page_complete."""
+    # Use a small PAGE_SIZE so we don't need thousands of results
+    small_page = 5
+    monkeypatch.setattr(unredact.app, "PAGE_SIZE", small_page)
+
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://test") as client:
+        resp = await client.get("/api/fonts")
+        fonts = resp.json()["fonts"]
+        font = next(f for f in fonts if f["available"])
+
+        # Word mode with small vocab but wide tolerance to generate many matches (>> 5)
+        resp = await client.post("/api/solve", json={
+            "font_id": font["id"],
+            "font_size": 40,
+            "gap_width_px": 100.0,
+            "tolerance_px": 20.0,
+            "left_context": "",
+            "right_context": "",
+            "hints": {"charset": "lowercase"},
+            "mode": "word",
+            "vocab_size": 200,
+        })
+        assert resp.status_code == 200
+        assert "text/event-stream" in resp.headers.get("content-type", "")
+
+        solve_id = resp.headers["x-solve-id"]
+        events = _parse_sse_events(resp.text)
+
+        # Separate event types
+        match_events = [e for e in events if e.get("status") == "match"]
+        page_complete_events = [e for e in events if e.get("status") == "page_complete"]
+        done_events = [e for e in events if e.get("status") == "done"]
+
+        # Exactly PAGE_SIZE match events should be streamed
+        assert len(match_events) == small_page, (
+            f"Expected {small_page} match events, got {len(match_events)}"
+        )
+
+        # A page_complete event should be present
+        assert len(page_complete_events) == 1
+        pc = page_complete_events[0]
+        assert pc["sent"] == small_page
+        assert pc["solve_id"] == solve_id
+
+        # The done event should still be present with total_found > PAGE_SIZE
+        assert len(done_events) == 1
+        total = done_events[0]["total_found"]
+        assert total > small_page, (
+            f"Expected total_found > {small_page}, got {total}"
+        )
+
+        # The buffer should contain ALL results (including the ones streamed)
+        assert solve_id in _solve_results
+        assert len(_solve_results[solve_id]) == total

unredact/app.py

@@ -370,10 +370,29 @@ async def solve(req: SolveRequest):
     _active_solves[solve_id] = False
 
     async def event_generator():
+        _solve_results.clear()
+        _solve_results[solve_id] = []
+
         try:
             found_texts = set()
+            sent_count = 0
+            page_complete_sent = False
             charset_name = req.hints.get("charset", "lowercase")
 
+            def _record(d):
+                """Buffer result and return True if it should be yielded as SSE."""
+                nonlocal sent_count, page_complete_sent
+                _solve_results[solve_id].append(d)
+                sent_count += 1
+                return sent_count <= PAGE_SIZE
+
+            def _page_complete_event():
+                nonlocal page_complete_sent
+                if not page_complete_sent:
+                    page_complete_sent = True
+                    return json.dumps({"status": "page_complete", "sent": PAGE_SIZE, "solve_id": solve_id})
+                return None
+
             # Name mode: single-word associate names
             if req.mode == "name" and not _active_solves.get(solve_id):
                 name_results = solve_name_dictionary(
@@ -389,13 +408,19 @@ async def solve(req: SolveRequest):
                     if r.text in found_texts:
                         continue
                     found_texts.add(r.text)
-                    yield json.dumps({
+                    d = {
                         "status": "match",
                         "text": r.text,
                         "width_px": round(r.width, 2),
                         "error_px": round(r.error, 2),
                         "source": "names",
-                    })
+                    }
+                    if _record(d):
+                        yield json.dumps(d)
+                    else:
+                        evt = _page_complete_event()
+                        if evt:
+                            yield evt
 
             # Full name mode: first x last Cartesian product + variants
             if req.mode == "full_name" and not _active_solves.get(solve_id):
@@ -412,13 +437,19 @@ async def solve(req: SolveRequest):
                     if r.text in found_texts:
                         continue
                     found_texts.add(r.text)
-                    yield json.dumps({
+                    d = {
                         "status": "match",
                         "text": r.text,
                         "width_px": round(r.width, 2),
                         "error_px": round(r.error, 2),
                         "source": "names",
-                    })
+                    }
+                    if _record(d):
+                        yield json.dumps(d)
+                    else:
+                        evt = _page_complete_event()
+                        if evt:
+                            yield evt
 
             # Email mode
             if req.mode == "email" and not _active_solves.get(solve_id):
@@ -432,13 +463,19 @@ async def solve(req: SolveRequest):
                         if _active_solves.get(solve_id):
                             break
                         found_texts.add(r.text)
-                        yield json.dumps({
+                        d = {
                             "status": "match",
                             "text": r.text,
                             "width_px": round(r.width, 2),
                             "error_px": round(r.error, 2),
                             "source": "emails",
-                        })
+                        }
+                        if _record(d):
+                            yield json.dumps(d)
+                        else:
+                            evt = _page_complete_event()
+                            if evt:
+                                yield evt
 
             # Word mode: English nouns + adjective/noun phrases
             if req.mode == "word" and not _active_solves.get(solve_id):
@@ -456,13 +493,19 @@ async def solve(req: SolveRequest):
                     if r.text in found_texts:
                         continue
                     found_texts.add(r.text)
-                    yield json.dumps({
+                    d = {
                         "status": "match",
                         "text": r.text,
                         "width_px": round(r.width, 2),
                         "error_px": round(r.error, 2),
                         "source": "words",
-                    })
+                    }
+                    if _record(d):
+                        yield json.dumps(d)
+                    else:
+                        evt = _page_complete_event()
+                        if evt:
+                            yield evt
 
             # Enumerate mode: Rust backend
             if req.mode == "enumerate" and not _active_solves.get(solve_id):
@@ -502,13 +545,19 @@ async def solve(req: SolveRequest):
                                 if text in found_texts:
                                     continue
                                 found_texts.add(text)
-                                yield json.dumps({
+                                d = {
                                     "status": "match",
                                     "text": text,
                                     "width_px": round(r["width"], 2),
                                     "error_px": round(r["error"], 2),
                                     "source": "enumerate",
-                                })
+                                }
+                                if _record(d):
+                                    yield json.dumps(d)
+                                else:
+                                    evt = _page_complete_event()
+                                    if evt:
+                                        yield evt
 
             yield json.dumps({
                 "status": "done",