fix(sql): enforce a real server-side row cap

agentkit/tools/builtin/sql_readonly.py

@@ -65,6 +65,15 @@
     re.IGNORECASE,
 )
 
+#: Regex matching a LIMIT clause with count, ALL, or NULL.
+_LIMIT_BOUND_REGEX = re.compile(r"\blimit\s+(\d+|all|null)\b", re.IGNORECASE)
+
+#: Regex matching a FETCH FIRST/NEXT clause.
+_FETCH_BOUND_REGEX = re.compile(
+    r"\bfetch\s+(first|next)\s*(?:(\d+|all|null)\s+)?(rows?)\s+(only|with\s+ties)\b",
+    re.IGNORECASE,
+)
+
 #: Process-wide engine, created on first use. Connection pools are expensive to rebuild and
 #: are safe to share, but the engine is still injectable per call so tests and alternative
 #: deployments can supply their own without mutating module state.
@@ -161,7 +170,7 @@ def _mask_literals(sql: str) -> str:
         char = sql[index]
 
         if char == "'":
-            masked.append(" ")
+            start = index
             index += 1
             while index < length:
                 if sql[index] == "'":
@@ -172,10 +181,11 @@ def _mask_literals(sql: str) -> str:
                     index += 1
                     break
                 index += 1
+            masked.append(" " * (index - start))
             continue
 
         if char == '"':
-            masked.append(" ")
+            start = index
             index += 1
             while index < length:
                 if sql[index] == '"':
@@ -185,6 +195,7 @@ def _mask_literals(sql: str) -> str:
                     index += 1
                     break
                 index += 1
+            masked.append(" " * (index - start))
             continue
 
         masked.append(char)
@@ -266,13 +277,8 @@ async def _apply_statement_timeout(conn: AsyncConnection, timeout_s: float) -> N
         logger.warning("Could not apply server-side statement timeout: %s", exc)
 
 
-def _top_level_tail(masked_sql: str) -> str:
-    """Return the trailing top-level clause of a masked statement.
-
-    Clause keywords are only meaningful at parenthesis depth zero: a LIMIT inside a subquery
-    does not bound the outer statement. Parenthesis counting on the literal-masked text is
-    exact for real queries, since quoted parentheses are already blanked out.
-    """
+def _top_level_tail_index(masked_sql: str) -> int:
+    """Return the start index of the trailing top-level clause in a masked statement."""
     depth = 0
     tail_start = 0
     for index, char in enumerate(masked_sql):
@@ -283,26 +289,87 @@ def _top_level_tail(masked_sql: str) -> str:
             if depth <= 0:
                 depth = 0
                 tail_start = index + 1
-    return masked_sql[tail_start:]
+    return tail_start
+
+
+def _top_level_tail(masked_sql: str) -> str:
+    """Return the trailing top-level clause of a masked statement.
+
+    Clause keywords are only meaningful at parenthesis depth zero: a LIMIT inside a subquery
+    does not bound the outer statement. Parenthesis counting on the literal-masked text is
+    exact for real queries, since quoted parentheses are already blanked out.
+    """
+    return masked_sql[_top_level_tail_index(masked_sql) :]
 
 
 def _apply_row_limit(sql: str, row_limit: int) -> str:
-    """Bound the statement's output in SQL, preferring an appended LIMIT.
+    """Bound the statement's output in SQL, preferring an appended or clamped LIMIT.
 
     Appending ``LIMIT`` leaves an already-valid statement untouched in shape, so constructs
-    like duplicate output column names in a JOIN keep working exactly as written. Only two
-    cases need different handling, both detected on the trailing top-level clause:
-
-    - the statement already bounds itself (LIMIT / FETCH FIRST): executed untouched;
-    - the statement ends in a row-locking clause (FOR UPDATE / FOR SHARE), after which a
-      LIMIT cannot legally appear: wrapped in a limited subquery instead.
+    like duplicate output column names in a JOIN keep working exactly as written.
+    Top-level bounds are handled as follows:
+
+    - if the statement already bounds itself (LIMIT / FETCH FIRST), caller-supplied bounds
+      larger than ``row_limit`` (or unbounded forms like LIMIT ALL / LIMIT NULL) are clamped
+      down to ``row_limit``, while smaller bounds are preserved;
+    - if the statement ends in a row-locking clause (FOR UPDATE / FOR SHARE) without a LIMIT,
+      after which a LIMIT cannot legally appear: wrapped in a limited subquery instead;
+    - otherwise: a LIMIT clause is appended directly.
     """
-    tail = _top_level_tail(_mask_literals(sql))
-    if _EXISTING_BOUND_REGEX.search(tail):
-        return sql
-    if _LOCKING_CLAUSE_REGEX.search(tail):
-        return f"SELECT * FROM ({sql}) AS agentkit_limited_query LIMIT {int(row_limit)}"
-    return f"{sql} LIMIT {int(row_limit)}"
+    row_limit = int(row_limit)
+    masked = _mask_literals(sql)
+    tail_start = _top_level_tail_index(masked)
+    head = sql[:tail_start]
+    tail = sql[tail_start:]
+    masked_tail = masked[tail_start:]
+
+    # Check for top-level LIMIT
+    limit_match = _LIMIT_BOUND_REGEX.search(masked_tail)
+    if limit_match:
+        val_str = limit_match.group(1).lower()
+        if val_str in ("all", "null"):
+            new_tail = (
+                tail[: limit_match.start()] + f"LIMIT {row_limit}" + tail[limit_match.end() :]
+            )
+            return head + new_tail
+        try:
+            val_int = int(val_str)
+            if val_int > row_limit:
+                new_tail = (
+                    tail[: limit_match.start()] + f"LIMIT {row_limit}" + tail[limit_match.end() :]
+                )
+                return head + new_tail
+            return sql
+        except ValueError:
+            return sql
+
+    # Check for top-level FETCH FIRST / NEXT
+    fetch_match = _FETCH_BOUND_REGEX.search(masked_tail)
+    if fetch_match:
+        direction = fetch_match.group(1)
+        count_str = fetch_match.group(2)
+        row_keyword = fetch_match.group(3)
+        tie_clause = fetch_match.group(4)
+
+        if count_str is None or count_str.lower() in ("all", "null"):
+            new_bound = f"FETCH {direction} {row_limit} {row_keyword} {tie_clause}"
+            new_tail = tail[: fetch_match.start()] + new_bound + tail[fetch_match.end() :]
+            return head + new_tail
+        try:
+            val_int = int(count_str)
+            if val_int > row_limit:
+                new_bound = f"FETCH {direction} {row_limit} {row_keyword} {tie_clause}"
+                new_tail = tail[: fetch_match.start()] + new_bound + tail[fetch_match.end() :]
+                return head + new_tail
+            return sql
+        except ValueError:
+            return sql
+
+    # Check for top-level locking clause (e.g. FOR UPDATE)
+    if _LOCKING_CLAUSE_REGEX.search(masked_tail):
+        return f"SELECT * FROM ({sql}) AS agentkit_limited_query LIMIT {row_limit}"
+
+    return f"{sql} LIMIT {row_limit}"
 
 
 @tool

docs/PROGRESS.md

@@ -407,3 +407,10 @@ This document tracks execution progress across all backlog tasks, including pull
 - **PR**: #47
 - **Commits**: 3
 - **Notes**: Remediated 17 findings from a full-codebase audit, rebased onto the AUD-01/AUD-04/AUD-05 follow-ups. Critical: removed the `ak_test_key_12345` default from `Settings.API_KEYS` and made the field required; replaced the `FakeLLMClient` fallback in `get_llm_client` with a `ServiceUnavailableError` (503) boundary around provider construction; kept the AUD-04 factory-level `RetryingLLMClient` wrapping and made it configurable through new `LLM_MAX_RETRIES`/`LLM_RETRY_BASE_DELAY_S`/`LLM_RETRY_MAX_DELAY_S` settings; made the per-request `Agent` assembly's override explicit as `app.state.agent_override`; forwarded `get_run_trace`/`list_runs` through `TieredMemoryStore`; moved `http_fetch` DNS off the event loop with `asyncio.to_thread`, added a post-connect `server_addr` check against the validated addresses to close the DNS-rebinding TOCTOU window, and replaced full-body buffering with a byte-budgeted stream; replaced the rate limiter's non-scripting fallback with `WATCH`/`MULTI`/`EXEC`, leaving the AUD-10 in-memory fallback for Redis outages intact. Moderate: corrected the `RedisMemoryStore.truncate_messages` off-by-one that returned `max_messages + 1` entries when a system prompt was present; switched generated `tool_call_id`s to full UUIDs; wrapped response parsing in both LLM adapters inside the error boundary so malformed responses and malformed tool-argument JSON raise `LLMProviderError`; guarded `RetryingLLMClient.__getattr__` against recursion on a partially constructed instance; sanitized `/health` details so the unauthenticated endpoint no longer echoes driver error strings; made `sql_readonly` mask string literals and quoted identifiers before its keyword scan, allow `#>`/`#>>`/`#-`, apply the row limit in SQL and push `statement_timeout` to PostgreSQL. Minor: stopped logging raw run goals at INFO (length only, content at DEBUG); replaced `split(";")` in `scripts/seed.py` with a statement-aware splitter handling literals, comments and dollar-quotes; bounded calculator exponentiation by result bit length. Added `tests/unit/test_audit_regressions.py` and a real-Redis concurrency regression test pinning the fallback fix. Recorded in ADR-009, ADR-010, ADR-011. Suite: 298 passed, 25 skipped; coverage 92.86% on core/tools/llm; `ruff check`, `ruff format --check` and `mypy agentkit tests` all clean.
+
+### AUDIT-02: Enforce Real Server-Side SQL Row Cap
+- **Status**: Done
+- **Branch**: `fix/sql-row-limit-cap`
+- **Issue**: #59
+- **Commits**: 1
+- **Notes**: In `agentkit/tools/builtin/sql_readonly.py`, `_apply_row_limit` now parses and clamps any top-level `LIMIT` (including `LIMIT ALL`, `LIMIT NULL`, and oversized limits) down to the configured `row_limit`, and clamps top-level `FETCH FIRST/NEXT` clauses. Preserves subqueries, valid smaller limits, and row-locking clauses (`FOR UPDATE`). Added comprehensive regression tests in `tests/unit/test_audit_regressions.py` verifying SQL text rewriting, server-side parameter passing, and literal masking.

tests/unit/test_audit_regressions.py

@@ -11,7 +11,7 @@
 from contextlib import asynccontextmanager
 from types import SimpleNamespace
 from typing import Any
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import fakeredis.aioredis
 import pytest
@@ -507,6 +507,116 @@ def test_row_limit_rewrite_prefers_appending_over_wrapping() -> None:
     assert "agentkit_limited_query" in wrapped and wrapped.endswith("LIMIT 50")
 
 
+def test_row_limit_clamps_oversized_and_unbounded_clauses() -> None:
+    """AUDIT-02: LIMIT ALL, LIMIT NULL, oversized limits and FETCH FIRST must be clamped."""
+    from agentkit.tools.builtin.sql_readonly import _apply_row_limit
+
+    # 1. LIMIT ALL and LIMIT NULL are replaced with configured row_limit
+    assert _apply_row_limit("SELECT * FROM runs LIMIT ALL", 50) == "SELECT * FROM runs LIMIT 50"
+    assert _apply_row_limit("SELECT * FROM runs LIMIT NULL", 50) == "SELECT * FROM runs LIMIT 50"
+
+    # 2. Oversized limits are clamped down to configured row_limit
+    assert (
+        _apply_row_limit("SELECT * FROM runs LIMIT 1000000000", 50) == "SELECT * FROM runs LIMIT 50"
+    )
+    assert _apply_row_limit("SELECT * FROM runs LIMIT 100", 50) == "SELECT * FROM runs LIMIT 50"
+
+    # 3. Smaller limits are preserved untouched
+    assert _apply_row_limit("SELECT * FROM runs LIMIT 5", 50) == "SELECT * FROM runs LIMIT 5"
+    assert _apply_row_limit("SELECT * FROM runs LIMIT 50", 50) == "SELECT * FROM runs LIMIT 50"
+
+    # 4. LIMIT with OFFSET
+    assert (
+        _apply_row_limit("SELECT * FROM runs LIMIT 100 OFFSET 10", 50)
+        == "SELECT * FROM runs LIMIT 50 OFFSET 10"
+    )
+    assert (
+        _apply_row_limit("SELECT * FROM runs OFFSET 10 LIMIT 100", 50)
+        == "SELECT * FROM runs OFFSET 10 LIMIT 50"
+    )
+
+    # 5. FETCH FIRST / NEXT clauses
+    assert (
+        _apply_row_limit("SELECT * FROM runs FETCH FIRST 100 ROWS ONLY", 50)
+        == "SELECT * FROM runs FETCH FIRST 50 ROWS ONLY"
+    )
+    assert (
+        _apply_row_limit("SELECT * FROM runs FETCH FIRST 5 ROWS ONLY", 50)
+        == "SELECT * FROM runs FETCH FIRST 5 ROWS ONLY"
+    )
+    assert (
+        _apply_row_limit("SELECT * FROM runs FETCH FIRST ROWS ONLY", 50)
+        == "SELECT * FROM runs FETCH FIRST 50 ROWS ONLY"
+    )
+    assert (
+        _apply_row_limit("SELECT * FROM runs FETCH NEXT 100 ROWS ONLY", 50)
+        == "SELECT * FROM runs FETCH NEXT 50 ROWS ONLY"
+    )
+
+    # 6. LIMIT inside a subquery is preserved; outer statement is bounded
+    subquery_no_outer_limit = "SELECT * FROM (SELECT * FROM t LIMIT 10) AS sub"
+    assert (
+        _apply_row_limit(subquery_no_outer_limit, 50)
+        == "SELECT * FROM (SELECT * FROM t LIMIT 10) AS sub LIMIT 50"
+    )
+
+    subquery_with_oversized_outer_limit = (
+        "SELECT * FROM (SELECT * FROM t LIMIT 10) AS sub LIMIT 100"
+    )
+    assert (
+        _apply_row_limit(subquery_with_oversized_outer_limit, 50)
+        == "SELECT * FROM (SELECT * FROM t LIMIT 10) AS sub LIMIT 50"
+    )
+
+    # 7. Locking clauses with LIMIT
+    assert (
+        _apply_row_limit("SELECT * FROM t LIMIT 100 FOR UPDATE", 50)
+        == "SELECT * FROM t LIMIT 50 FOR UPDATE"
+    )
+    assert (
+        _apply_row_limit("SELECT * FROM t LIMIT 5 FOR UPDATE", 50)
+        == "SELECT * FROM t LIMIT 5 FOR UPDATE"
+    )
+    assert (
+        _apply_row_limit("SELECT * FROM t LIMIT ALL FOR UPDATE", 50)
+        == "SELECT * FROM t LIMIT 50 FOR UPDATE"
+    )
+
+    # 8. Literals containing the word 'limit' must not be confused with SQL clause
+    literal_query = "SELECT * FROM t WHERE note = 'limit all' LIMIT 100"
+    assert (
+        _apply_row_limit(literal_query, 50) == "SELECT * FROM t WHERE note = 'limit all' LIMIT 50"
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_sql_query_enforces_limit_server_side() -> None:
+    """AUDIT-02: verify conn.execute() receives the clamped LIMIT in SQL."""
+    from agentkit.tools.builtin.sql_readonly import execute_sql_query
+
+    mock_engine = AsyncMock()
+    mock_conn = AsyncMock()
+    mock_cursor = MagicMock()
+    mock_cursor.keys.return_value = ["id", "val"]
+    mock_cursor.fetchmany.return_value = [(1, "test")]
+    mock_conn.execute.return_value = mock_cursor
+
+    @asynccontextmanager
+    async def _mock_connect() -> AsyncIterator[Any]:
+        yield mock_conn
+
+    mock_engine.connect = _mock_connect
+
+    await execute_sql_query("SELECT * FROM runs LIMIT ALL", engine=mock_engine, row_limit=50)
+
+    # Inspect the exact SQL text passed to conn.execute
+    executed_statements = [str(call.args[0]) for call in mock_conn.execute.call_args_list]
+    query_statements = [s for s in executed_statements if s.startswith("SELECT")]
+    assert len(query_statements) == 1
+    assert query_statements[0] == "SELECT * FROM runs LIMIT 50"
+    assert "LIMIT ALL" not in query_statements[0]
+
+
 # --- seed script: statement-aware splitting -------------------------------------------------