diff --git a/agentkit/tools/builtin/sql_readonly.py b/agentkit/tools/builtin/sql_readonly.py index fa3e905..de11407 100644 --- a/agentkit/tools/builtin/sql_readonly.py +++ b/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 diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 2c7b115..e04bf6a 100644 --- a/docs/PROGRESS.md +++ b/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. diff --git a/tests/unit/test_audit_regressions.py b/tests/unit/test_audit_regressions.py index e7c0d9d..8dcfba9 100644 --- a/tests/unit/test_audit_regressions.py +++ b/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 -------------------------------------------------