From 6bb655f63668952dd0e2f7606915dfb0604f6bca Mon Sep 17 00:00:00 2001
From: Rudra Prasad Bhuyan <rudraprasadbhuyan999@gmail.com>
Date: Wed, 16 Sep 2026 11:06:33 +0000
Subject: [PATCH] fix: Search limit validation

---
 tests/test_acl.py     | 30 +++++++++++++++++++++++++++++-
 vaultrag/retrieval.py |  5 +++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/tests/test_acl.py b/tests/test_acl.py
index ebad5bc..b67b5ac 100644
--- a/tests/test_acl.py
+++ b/tests/test_acl.py
@@ -16,7 +16,7 @@
 
 import pytest
 
-from vaultrag.retrieval import resolve_principal, search
+from vaultrag.retrieval import Principal, resolve_principal, search
 
 pytestmark = pytest.mark.asyncio
 
@@ -31,6 +31,34 @@ async def _search_as(conn, embedder, user_id: str, question: str = QUESTION, lim
     return {h.doc_id for h in hits}
 
 
+class _NoDatabaseAccess:
+    def cursor(self, **kwargs):
+        raise AssertionError("search opened a database cursor before validating arguments")
+
+
+@pytest.mark.parametrize(
+    ("limit", "candidates", "message"),
+    [
+        (0, 50, "limit"),
+        (-1, 50, "limit"),
+        (5, 0, "candidates"),
+        (5, -1, "candidates"),
+    ],
+)
+async def test_search_rejects_non_positive_counts_before_database_access(
+    limit, candidates, message
+):
+    with pytest.raises(ValueError, match=message):
+        await search(
+            _NoDatabaseAccess(),
+            Principal(user_id="alice", groups=()),
+            QUESTION,
+            [0.0],
+            limit=limit,
+            candidates=candidates,
+        )
+
+
 async def test_group_isolation_alice_cannot_see_sales_or_hr(conn, corpus, embedder):
     docs = await _search_as(conn, embedder, "alice")
     assert "eng-handbook" in docs, "alice is in engineering and must see the eng handbook"
diff --git a/vaultrag/retrieval.py b/vaultrag/retrieval.py
index 76e671a..dc458ce 100644
--- a/vaultrag/retrieval.py
+++ b/vaultrag/retrieval.py
@@ -77,6 +77,11 @@ async def search(
     Returns:
         Up to `limit` chunks the user is allowed to see, best first.
     """
+    if limit <= 0:
+        raise ValueError("limit must be greater than 0")
+    if candidates <= 0:
+        raise ValueError("candidates must be greater than 0")
+
     # A chunk is visible iff its document has an ACL row whose principal is one of ours.
     # EXISTS rather than JOIN so a document with several matching ACL rows yields one chunk, not N.
     #