From 88bea9b24c494586d28a586496e07a4b6d7fe35c Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 15:06:00 -0700 Subject: [PATCH 01/10] feat(executor): the timeout resolver and the deadline primitive (S1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the config surface and the watchdog that a per-statement timeout will sit on, with nothing wired into an engine yet. `_resolve_timeout_s` mirrors `_resolve_row_cap`: `AGAMI_SQL_TIMEOUT_S` is the operator-configurable deployment budget, 30s by default, and a missing or non-positive value falls back. It differs in one deliberate way — a value that is present but unparseable is logged at warning naming the rejected text before falling back, rather than being swallowed, because an operator who wrote `45.5` asked for something specific. The warning goes to the module logger and never to stderr, which the subprocess transport parses. `_deadline` is a context manager that arms a daemon `threading.Timer`, sets its Event before invoking the driver's cancel so the flag is observable by whoever catches the resulting error, swallows and logs a cancel that raises from the timer thread, and disarms in a `finally`. `_ResourceLimit` is the internal marker a cancelled statement will unwind on. Spec: ACE-038 Co-Authored-By: Claude Opus 5 (1M context) --- packages/agami-core/src/execute_sql.py | 75 ++++++++- plugins/agami/lib/execute_sql.py | 75 ++++++++- tests/test_ace038_timeout.py | 220 +++++++++++++++++++++++++ 3 files changed, 368 insertions(+), 2 deletions(-) create mode 100644 tests/test_ace038_timeout.py diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 9a5b51d..3d8efff 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -52,15 +52,17 @@ import argparse import configparser +import contextlib import csv import json import logging import os import stat import sys +import threading import urllib.parse import uuid -from collections.abc import Callable +from collections.abc import Callable, Iterator from contextvars import ContextVar from dataclasses import asdict, dataclass from pathlib import Path @@ -860,6 +862,77 @@ def _resolve_row_cap() -> int: return cap +_DEFAULT_TIMEOUT_S = 30 # wall-clock seconds one statement may run before the watchdog cancels it +# Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the +# HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the +# context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one +# thread) it behaves exactly as a module global would. +_timeout_override: ContextVar[int | None] = ContextVar("_timeout_override", default=None) + + +def _resolve_timeout_s() -> int: + """Effective per-statement timeout, in whole seconds. `AGAMI_SQL_TIMEOUT_S` is the + operator-configurable DEPLOYMENT budget (default 30 when unset) — an operator owns their + availability tradeoff and may set it higher OR lower than 30. A missing or non-positive value + falls back to the default. + + Unlike `_resolve_row_cap`, a value that is PRESENT but unparseable is logged at warning before + the fallback. An operator who wrote `45.5` or `30s` asked for something specific and silently + running 30 instead is how a misconfiguration survives a whole deployment unnoticed. The warning + goes to the module logger and never to stderr, because the subprocess transport parses stderr and + an extra line there would break that contract.""" + raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() + digits = raw[1:] if raw.startswith("-") else raw # a leading minus is a value, not a typo + if raw and not digits.isdigit(): + _LOG.warning( + "AGAMI_SQL_TIMEOUT_S=%r is not a whole number of seconds; falling back to %ds.", + raw, + _DEFAULT_TIMEOUT_S, + ) + timeout_s = int(raw) if digits.isdigit() else _DEFAULT_TIMEOUT_S + if timeout_s <= 0: + timeout_s = _DEFAULT_TIMEOUT_S # "0" / "-5" → the default, never an instantly-expired budget + override = _timeout_override.get() + if override is not None and override > 0: + timeout_s = override # a caller that resolved its own budget outranks the deployment default + return timeout_s + + +class _ResourceLimit(Exception): + """Raised when our own watchdog fired, so a cancelled statement unwinds the engine function the + same way any other failure does and the surrounding transaction rolls back. It is an internal + marker only: it is always caught and translated inside this module and never crosses the tool + boundary.""" + + +@contextlib.contextmanager +def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threading.Event]: + """Arm a watchdog that calls `cancel` if the wrapped block outlives `timeout_s`, and yield the + `threading.Event` that says whether it fired. + + The event is set BEFORE `cancel` runs. Order matters: whoever catches the driver error that the + cancellation provokes must be able to read an already-set flag and attribute the failure to us + rather than to the database. A `cancel` that raises is swallowed and logged, because some drivers + raise when cancelled from a thread other than the one running the statement, and an exception + escaping a timer thread is both unhandleable by the caller and invisible in the result.""" + fired = threading.Event() + + def fire() -> None: + fired.set() + try: + cancel() + except Exception as exc: + _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) + + timer = threading.Timer(timeout_s, fire) + timer.daemon = True # a hung cancel must never hold the interpreter open at shutdown + timer.start() + try: + yield fired + finally: + timer.cancel() # a block that finished on time disarms the watchdog before it can fire + + def _flag_truncated(cap: int) -> None: """Signal a bounded-fetch truncation to the caller — a non-error `{"truncated": …}` marker on stderr (distinct from the guards' `{"error": …}`), so a truncated result is never mistaken for a diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 9a5b51d..3d8efff 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -52,15 +52,17 @@ import argparse import configparser +import contextlib import csv import json import logging import os import stat import sys +import threading import urllib.parse import uuid -from collections.abc import Callable +from collections.abc import Callable, Iterator from contextvars import ContextVar from dataclasses import asdict, dataclass from pathlib import Path @@ -860,6 +862,77 @@ def _resolve_row_cap() -> int: return cap +_DEFAULT_TIMEOUT_S = 30 # wall-clock seconds one statement may run before the watchdog cancels it +# Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the +# HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the +# context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one +# thread) it behaves exactly as a module global would. +_timeout_override: ContextVar[int | None] = ContextVar("_timeout_override", default=None) + + +def _resolve_timeout_s() -> int: + """Effective per-statement timeout, in whole seconds. `AGAMI_SQL_TIMEOUT_S` is the + operator-configurable DEPLOYMENT budget (default 30 when unset) — an operator owns their + availability tradeoff and may set it higher OR lower than 30. A missing or non-positive value + falls back to the default. + + Unlike `_resolve_row_cap`, a value that is PRESENT but unparseable is logged at warning before + the fallback. An operator who wrote `45.5` or `30s` asked for something specific and silently + running 30 instead is how a misconfiguration survives a whole deployment unnoticed. The warning + goes to the module logger and never to stderr, because the subprocess transport parses stderr and + an extra line there would break that contract.""" + raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() + digits = raw[1:] if raw.startswith("-") else raw # a leading minus is a value, not a typo + if raw and not digits.isdigit(): + _LOG.warning( + "AGAMI_SQL_TIMEOUT_S=%r is not a whole number of seconds; falling back to %ds.", + raw, + _DEFAULT_TIMEOUT_S, + ) + timeout_s = int(raw) if digits.isdigit() else _DEFAULT_TIMEOUT_S + if timeout_s <= 0: + timeout_s = _DEFAULT_TIMEOUT_S # "0" / "-5" → the default, never an instantly-expired budget + override = _timeout_override.get() + if override is not None and override > 0: + timeout_s = override # a caller that resolved its own budget outranks the deployment default + return timeout_s + + +class _ResourceLimit(Exception): + """Raised when our own watchdog fired, so a cancelled statement unwinds the engine function the + same way any other failure does and the surrounding transaction rolls back. It is an internal + marker only: it is always caught and translated inside this module and never crosses the tool + boundary.""" + + +@contextlib.contextmanager +def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threading.Event]: + """Arm a watchdog that calls `cancel` if the wrapped block outlives `timeout_s`, and yield the + `threading.Event` that says whether it fired. + + The event is set BEFORE `cancel` runs. Order matters: whoever catches the driver error that the + cancellation provokes must be able to read an already-set flag and attribute the failure to us + rather than to the database. A `cancel` that raises is swallowed and logged, because some drivers + raise when cancelled from a thread other than the one running the statement, and an exception + escaping a timer thread is both unhandleable by the caller and invisible in the result.""" + fired = threading.Event() + + def fire() -> None: + fired.set() + try: + cancel() + except Exception as exc: + _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) + + timer = threading.Timer(timeout_s, fire) + timer.daemon = True # a hung cancel must never hold the interpreter open at shutdown + timer.start() + try: + yield fired + finally: + timer.cancel() # a block that finished on time disarms the watchdog before it can fire + + def _flag_truncated(cap: int) -> None: """Signal a bounded-fetch truncation to the caller — a non-error `{"truncated": …}` marker on stderr (distinct from the guards' `{"error": …}`), so a truncated result is never mistaken for a diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py new file mode 100644 index 0000000..1312685 --- /dev/null +++ b/tests/test_ace038_timeout.py @@ -0,0 +1,220 @@ +"""Per-statement timeout — config resolution and the deadline primitive. + +`_resolve_timeout_s` answers "how long may one statement run", from `AGAMI_SQL_TIMEOUT_S`, the +`_timeout_override` ContextVar, and a 30s default; unlike the row cap it complains out loud when the +env value is present but unparseable. `_deadline` is the watchdog those seconds feed: it fires an +Event and calls a cancel callable when a block outlives its budget, and disarms cleanly when it does +not. Nothing in the engines calls either yet. +""" + +from __future__ import annotations + +import logging +import sys +import threading +import time +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +PKG_SRC = REPO_ROOT / "packages" / "agami-core" / "src" +if str(PKG_SRC) not in sys.path: + sys.path.insert(0, str(PKG_SRC)) + +import execute_sql # noqa: E402 + +# Short enough that the whole file stays sub-second, long enough that a loaded CI runner still +# schedules the timer thread before the assertion runs. +_TINY = 0.05 + + +@pytest.fixture(autouse=True) +def _reset_override(): + # _timeout_override is a request-scoped ContextVar; isolate every test from it. + execute_sql._timeout_override.set(None) + yield + execute_sql._timeout_override.set(None) + + +@pytest.fixture(autouse=True) +def _clear_env(monkeypatch): + # The suite must not inherit an operator's real budget from the ambient environment. + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) + + +# -------------------------------------------------------------------------------------------- +# _resolve_timeout_s +# -------------------------------------------------------------------------------------------- + + +def test_an_absent_env_var_yields_the_default_budget(): + assert execute_sql._resolve_timeout_s() == 30 + assert execute_sql._DEFAULT_TIMEOUT_S == 30 + + +@pytest.mark.parametrize("raw", ["1", "5", "45", "600"]) +def test_a_valid_env_value_is_honoured(monkeypatch, raw): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + assert execute_sql._resolve_timeout_s() == int(raw) + + +def test_surrounding_whitespace_does_not_defeat_a_valid_value(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", " 45 ") + assert execute_sql._resolve_timeout_s() == 45 + + +@pytest.mark.parametrize("raw", ["0", "00", "-5"]) +def test_a_non_positive_value_falls_back_to_the_default(monkeypatch, raw): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + assert execute_sql._resolve_timeout_s() == 30 + + +@pytest.mark.parametrize("raw", ["6O", "30s", "45.5", "thirty", "1e3"]) +def test_an_unparseable_value_falls_back_and_says_so(monkeypatch, caplog, raw): + """The row cap falls back silently; this one must not. An operator who typed a capital O for a + zero has to be able to find out why their budget is not what they configured.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + assert execute_sql._resolve_timeout_s() == 30 + + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert warnings, f"no warning emitted for the rejected value {raw!r}" + assert any(raw in r.getMessage() for r in warnings), ( + f"the warning must name the rejected text {raw!r}; got {[r.getMessage() for r in warnings]}" + ) + + +@pytest.mark.parametrize("raw", ["", "45", "0", "-5"]) +def test_a_parseable_or_absent_value_stays_quiet(monkeypatch, caplog, raw): + """Only genuinely unparseable text warrants the warning. A deliberate `0` or `-5` is a value we + understood and declined, not a typo, and warning on it would train operators to ignore the log.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + execute_sql._resolve_timeout_s() + assert [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] == [] + + +def test_the_context_var_takes_precedence_over_the_env(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") + execute_sql._timeout_override.set(7) + assert execute_sql._resolve_timeout_s() == 7 + + +def test_the_context_var_takes_precedence_over_the_default(): + execute_sql._timeout_override.set(12) + assert execute_sql._resolve_timeout_s() == 12 + + +def test_the_context_var_may_raise_the_budget_as_well_as_lower_it(monkeypatch): + """Unlike the row cap, which can only be tightened per call, the override wins outright.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "5") + execute_sql._timeout_override.set(90) + assert execute_sql._resolve_timeout_s() == 90 + + +@pytest.mark.parametrize("override", [0, -1]) +def test_a_non_positive_override_is_ignored(monkeypatch, override): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") + execute_sql._timeout_override.set(override) + assert execute_sql._resolve_timeout_s() == 45 + + +# -------------------------------------------------------------------------------------------- +# _deadline +# -------------------------------------------------------------------------------------------- + + +class _RecordingCancel: + """A stand-in for a driver's `cancel()`, recording the calls and the order relative to the flag.""" + + def __init__(self, fails: bool = False): + self.calls = 0 + self.fails = fails + self.done = threading.Event() + + def __call__(self) -> None: + self.calls += 1 + try: + if self.fails: + raise RuntimeError("driver refused to cancel from another thread") + finally: + self.done.set() + + +def test_an_overrunning_block_sets_the_event_and_cancels(): + cancel = _RecordingCancel() + with execute_sql._deadline(cancel, _TINY) as fired: + assert cancel.done.wait(2.0), "the watchdog never ran" + # The flag must already be readable by the time the cancel lands, so a caller catching the + # resulting driver error can attribute it to us rather than to the database. + assert fired.is_set() + assert cancel.calls == 1 + + +def test_a_block_that_finishes_first_neither_fires_nor_cancels(): + cancel = _RecordingCancel() + with execute_sql._deadline(cancel, 30) as fired: + pass + assert not fired.is_set() + assert cancel.calls == 0 + + +def test_the_timer_is_disarmed_on_exit_so_no_late_cancel_arrives(): + """The watchdog must not outlive its block: a cancel landing after the statement finished would + kill whatever the connection is doing next.""" + cancel = _RecordingCancel() + with execute_sql._deadline(cancel, _TINY) as fired: + pass + time.sleep(_TINY * 4) # comfortably past when an un-disarmed timer would have fired + assert cancel.calls == 0 + assert not fired.is_set() + + +def test_the_timer_is_disarmed_even_when_the_block_raises(): + cancel = _RecordingCancel() + with pytest.raises(ValueError): + with execute_sql._deadline(cancel, _TINY): + raise ValueError("the statement blew up on its own") + time.sleep(_TINY * 4) + assert cancel.calls == 0 + + +def test_a_cancel_that_raises_does_not_escape_the_timer_thread(caplog): + """Some drivers raise when cancelled from a thread other than the one running the statement. That + must be logged and swallowed: an exception escaping a timer thread is unhandleable by the caller + and would be lost to threading's excepthook.""" + cancel = _RecordingCancel(fails=True) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + with execute_sql._deadline(cancel, _TINY) as fired: + assert cancel.done.wait(2.0), "the watchdog never ran" + time.sleep(_TINY) # let `fire` finish handling the raised cancel before we leave + + assert fired.is_set() # the timeout still counts as fired even though the cancel failed + assert cancel.calls == 1 + assert any( + r.levelno == logging.WARNING and "driver refused to cancel" in r.getMessage() + for r in caplog.records + ), f"the failed cancel was not logged; got {[r.getMessage() for r in caplog.records]}" + + +def test_the_watchdog_thread_is_a_daemon_and_does_not_hold_the_process_open(): + """A hung cancel must never keep the interpreter alive at shutdown.""" + live_before = {t.ident for t in threading.enumerate()} + with execute_sql._deadline(_RecordingCancel(), 300) as fired: + new = [t for t in threading.enumerate() if t.ident not in live_before] + assert new, "no watchdog thread was started" + assert all(t.daemon for t in new) + assert not fired.is_set() + + +# -------------------------------------------------------------------------------------------- +# _ResourceLimit +# -------------------------------------------------------------------------------------------- + + +def test_the_resource_limit_marker_is_a_plain_exception(): + """It has to unwind an engine function like any other error so the transaction rolls back.""" + assert issubclass(execute_sql._ResourceLimit, Exception) + with pytest.raises(execute_sql._ResourceLimit): + raise execute_sql._ResourceLimit("statement exceeded its budget") From fc272c4977dca37c712e2da0033bb6a68b54a406 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 15:31:09 -0700 Subject: [PATCH 02/10] feat(executor): the per-statement deadline, proven end to end on SQLite (S2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires S1's watchdog into one engine and turns a cancelled statement into the refusal the contract has been reserving `resource_limit` for. SQLite goes first because it is in-process, needs no network, and `Connection.interrupt()` is a genuine cancel rather than a request the driver may ignore. `_run_sqlite` resolves its budget once and wraps BOTH `cur.execute` and `_collect_cursor` in `_deadline(conn.interrupt, timeout_s)`. The fetch is inside the clock deliberately: `_collect_cursor` pulls `cap + 1` rows in a single `fetchmany`, and on a cursor that streams its result that pull is where the scan happens, so a clock stopping at `execute` would bound the cheap half. **The flag is the classification, and only the flag.** A cancelled statement raises `OperationalError("interrupted")`, so neither the message nor the elapsed clock can be the test — both are things an ordinary database error can have by coincidence, and reading either would tell an unlucky caller to narrow a query that never ran long. `_deadline` sets its Event before the cancel lands, so "did WE stop this?" is answered without inference. Asserted in both directions: the same error text with the flag clear stays `failed`, including when it arrives well past the budget. An `except _ResourceLimit: raise` sits ahead of `_run_sqlite`'s `code=5` catch-all, and ahead of both handlers in `execute_guarded`, so the marker reaches the chokepoint instead of being reported as the database's outcome. The other nine engines are unchanged and remain bounded only by the supervisor. The refusal quotes the configured budget — a deployment setting, not a data value — and names no environment variable: on the served path the caller is an assistant with no shell and no deployment, so "raise AGAMI_SQL_TIMEOUT_S" is advice aimed past it at an operator who is not in the conversation. Contract changes to existing tests, each because the contract changed: * `test_ace035_no_enumeration.py` — `resource_limit` had a `_NO_VECTOR` entry saying "the day something emits it, it needs a vector here". Today is that day: the entry is gone and a real vector takes its place, a runaway CTE that passes every gate and is stopped inside the executor. It is the only row in the matrix whose refusal is produced AFTER the model has been loaded and consulted, which is the state in which a helpful message has the declared surface closest to hand. The matrix is 20 rows, not 16. * `test_ace035_guardrail_audit.py` — `test_no_gate_produces_the_resource_limit_refusal` is renamed to `test_the_supervisor_kill_does_not_borrow_the_resource_limit_refusal`. Not weakened and not deleted: it drives the same branch and makes the same assertions. Only its premise moved — "no gate produces it" is no longer true, while "this branch must not borrow it" is now the sharper claim. * `guardrail.py` and `tools.py` docstrings asserting "nothing imposes one yet" now describe the producer. Spec: ACE-038 Co-Authored-By: Claude Opus 5 (1M context) --- packages/agami-core/src/execute_sql.py | 64 +++- packages/agami-core/src/guardrail.py | 25 +- packages/agami-core/src/tools.py | 6 +- plugins/agami/lib/execute_sql.py | 64 +++- plugins/agami/lib/guardrail.py | 25 +- tests/test_ace035_guardrail_audit.py | 21 +- tests/test_ace035_no_enumeration.py | 68 ++-- tests/test_ace038_timeout.py | 428 ++++++++++++++++++++++++- 8 files changed, 636 insertions(+), 65 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 3d8efff..62c6b80 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -72,6 +72,7 @@ from guardrail import ( RULE_MODEL_SAFETY, RULE_MODEL_UNAVAILABLE, + RULE_RESOURCE_LIMIT, Envelope, Failure, FailureKind, @@ -674,6 +675,12 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: + """Tier-3 path for SQLite using the stdlib `sqlite3` module. + + The first engine to run under the per-statement deadline. `sqlite3.Connection.interrupt()` is a + genuine cancel — it stops the statement mid-scan from another thread — so this engine can prove + the whole refusal contract in-process, with no network and no fixture warehouse. + """ import sqlite3 # always available in stdlib _require(creds, "path") path = os.path.expanduser(creds["path"]) @@ -681,10 +688,38 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: conn = sqlite3.connect(path) except Exception as e: raise ExecutorError(f"SQLite connect failed: {e}", code=4) + # Resolved ONCE for the call, so the budget the watchdog enforces and the number the refusal + # quotes cannot be two different values. + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The deadline covers the FETCH as well as the execute. `_collect_cursor` pulls `cap + 1` + # rows in a single `fetchmany`, and on a cursor that streams its result that pull is where + # the scan actually happens — a clock that stopped at `execute` would bound the cheap half + # of the work and leave the expensive half unbounded. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + # A cancelled statement raises `sqlite3.OperationalError("interrupted")`, whose text + # is not a classification we would want to key on. The FLAG is the classification, + # and it is the ONLY one: `_deadline` sets it before the cancel lands, so an error + # raised while it is unset is the database's own — however late it arrives. + if expired.is_set(): + raise _ResourceLimit("the statement outlived its per-statement budget") + raise + # Checked after the watchdog is disarmed, so the flag is final. A cancel can also land + # between the two calls above, or just as the second returns, and leave nothing to raise: + # the budget still elapsed, so the outcome is still a refusal rather than a result gathered + # past it. + if expired.is_set(): + raise _ResourceLimit("the statement outlived its per-statement budget") + except _ResourceLimit: + # Ahead of the catch-all below on purpose: our own marker must reach `execute_guarded` + # intact. Wrapped in an `ExecutorError` it would become a `failed`/`syntax` envelope, which + # is the database's outcome rather than the bound we imposed. + raise except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: @@ -1342,11 +1377,12 @@ def execute_guarded( **Every path returns exactly one ``Envelope``, and this function is TOTAL** — nothing is raised out of here for a caller to interpret, so the subprocess ``main`` and the in-process MCP handler - cannot disagree about what happened. The six outcomes: + cannot disagree about what happened. The seven outcomes: * the read-only gate refuses -> ``refused`` carrying that gate's ``Refusal`` * ``_model_safety`` returns a Refusal -> ``refused`` carrying it verbatim * ``_model_safety`` returns an int -> ``refused`` carrying the interim ``model_safety`` + * the per-statement deadline fired -> ``refused`` carrying ``resource_limit`` * ``executor.execute`` raises -> ``failed`` carrying a classified ``Failure`` * anything else raises -> ``failed``/``other``, generic message, raw to the log * the statement ran -> ``ok`` carrying the ``ExecResult`` @@ -1400,6 +1436,28 @@ def execute_guarded( # does not accept) fails the present-iff check in `Envelope.__post_init__`, and that is a # broken adapter, not a reason for the chokepoint to raise at its caller. return _envelope("ok", data=result) + except _ResourceLimit: + # AHEAD of both handlers below: `_ResourceLimit` is an `ExecutorError` sibling, not a + # subclass, but the catch-all would swallow it and report a bound WE imposed as an + # unclassified server break. This is the per-statement bound the contract reserves + # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a + # fix we can honestly name, unlike the supervisor's kill of a child that never returned. + # + # The budget is re-resolved rather than carried on the marker: the marker stays a plain + # exception, and nothing between the engine call and here can change the env var or the + # request-scoped ContextVar the resolver reads, so it is the same number the watchdog used. + timeout_s = _resolve_timeout_s() + return _envelope("refused", refusal=refuse( + RULE_RESOURCE_LIMIT, + # The configured number belongs here — it is a deployment setting, not a data value, and + # a bound the caller cannot see is one it cannot plan around. The remediation names only + # what would make THIS statement executable: on the served path the caller is an + # assistant with no shell and no deployment, so naming the environment variable would be + # advice it cannot take, addressed to someone who is not reading. + detail=f"The statement ran longer than the {timeout_s}s limit and was cancelled.", + remediation="Narrow the time range, reduce the grouping, or add a selective filter, " + "then run it again.", + )) except ExecutorError as exc: # The classified branch: `msg` is authored by this module (a missing driver, a connect # failure, the credential-resolution remediation naming DATASOURCE_URL) or relayed from the diff --git a/packages/agami-core/src/guardrail.py b/packages/agami-core/src/guardrail.py index 0e7e437..8c0bc07 100644 --- a/packages/agami-core/src/guardrail.py +++ b/packages/agami-core/src/guardrail.py @@ -50,11 +50,13 @@ RULE_COLUMN_SCOPE = "column_scope" RULE_SELECT_STAR = "select_star" RULE_MODEL_UNAVAILABLE = "model_unavailable" -# Declared and pinned below, with NO producer today. The contract reserves it for a **per-statement** -# timeout the guard imposes — a bound whose subject is the statement, so "narrow the query" is a fix -# we can honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this -# rule: it stops a child that never returned, without knowing what the child was doing when it -# stopped, so it is a `failed`/`timeout` (see `FailureKind` below and contract §3). +# Produced by the **per-statement** timeout the executor imposes: a watchdog cancels a statement that +# outlives the configured budget, and `execute_sql.execute_guarded` turns that into this refusal. Its +# subject is the statement, which is what earns it a rule at all — "narrow the query" is a fix we can +# honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this rule: it +# stops a child that never returned, without knowing what the child was doing when it stopped, so it +# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Wired into SQLite first; the +# other engines follow, and until they do a statement on them is bounded only by that supervisor. RULE_RESOURCE_LIMIT = "resource_limit" RULE_UNPARSEABLE = "unparseable" @@ -90,8 +92,8 @@ RULE_SELECT_STAR: "out_of_scope", RULE_MODEL_UNAVAILABLE: "undetermined", # A bound we imposed, not a property of the statement: neither unsafe nor out of scope — we - # simply did not determine the answer within the bound. Pinned here while unproduced so the gate - # that eventually imposes a per-statement timeout fills a constant rather than inventing one. + # simply did not determine the answer within the bound. Pinned before it had a producer, so the + # gate that imposes the per-statement timeout filled a constant rather than inventing one. RULE_RESOURCE_LIMIT: "undetermined", RULE_UNPARSEABLE: "undetermined", RULE_MODEL_SAFETY: "undetermined", @@ -165,9 +167,12 @@ def refuse(rule: str, *, detail: str, remediation: str) -> Refusal: in connect, credential resolution or model load, where "narrow the query" is the wrong fix. So an unresponsive executor is `failed` / `timeout`, and its message says only that we stopped waiting (guardrail contract §3). A **per-statement** timeout is the other case — its subject IS the -statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT` — and nothing imposes one yet, which is -why that rule is pinned with no producer. Driver-level connect/login timeouts fold into the connect -failure the executor already reports as `auth` (exit 4), because that is what the driver raises. +statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT`, and the executor now imposes one: a +watchdog cancels the statement through the driver and the outcome leaves the chokepoint on the +refusal channel rather than this one. The two therefore coexist and stay distinguishable in the audit +trail, which is the whole point of splitting them. Driver-level connect/login timeouts fold into the +connect failure the executor already reports as `auth` (exit 4), because that is what the driver +raises. `column_not_found`, `table_not_found`, `permission` and `network` are DECLARED BUT UNREACHABLE: producing them means parsing driver text, and sanitizing driver text belongs to the error-hardening diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index 36ba72f..1bbbd76 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -1428,8 +1428,10 @@ def tool_execute_sql(args: dict[str, Any]) -> str: # back; it does not say the STATEMENT ran long. The child may have hung in connect, in # credential resolution or in loading the model, and on any of those "narrow the query" is # advice pointing at the wrong thing. A refusal must name a fix, so a kill we cannot attribute - # is not one. (Guardrail contract §3: an unresponsive executor is `failed`. The per-statement - # timeout a later gate imposes IS a refusal, and `RULE_RESOURCE_LIMIT` is reserved for it.) + # is not one. (Guardrail contract §3: an unresponsive executor is `failed`. The executor's own + # per-statement deadline IS a refusal and carries `RULE_RESOURCE_LIMIT`; it can attribute the + # cancel to the statement because it is the thing it cancelled. The two bounds coexist, and + # this one must not borrow the other's rule.) return _emit( _envelope("failed", failure=Failure( kind="timeout", diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 3d8efff..62c6b80 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -72,6 +72,7 @@ from guardrail import ( RULE_MODEL_SAFETY, RULE_MODEL_UNAVAILABLE, + RULE_RESOURCE_LIMIT, Envelope, Failure, FailureKind, @@ -674,6 +675,12 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: + """Tier-3 path for SQLite using the stdlib `sqlite3` module. + + The first engine to run under the per-statement deadline. `sqlite3.Connection.interrupt()` is a + genuine cancel — it stops the statement mid-scan from another thread — so this engine can prove + the whole refusal contract in-process, with no network and no fixture warehouse. + """ import sqlite3 # always available in stdlib _require(creds, "path") path = os.path.expanduser(creds["path"]) @@ -681,10 +688,38 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: conn = sqlite3.connect(path) except Exception as e: raise ExecutorError(f"SQLite connect failed: {e}", code=4) + # Resolved ONCE for the call, so the budget the watchdog enforces and the number the refusal + # quotes cannot be two different values. + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The deadline covers the FETCH as well as the execute. `_collect_cursor` pulls `cap + 1` + # rows in a single `fetchmany`, and on a cursor that streams its result that pull is where + # the scan actually happens — a clock that stopped at `execute` would bound the cheap half + # of the work and leave the expensive half unbounded. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + # A cancelled statement raises `sqlite3.OperationalError("interrupted")`, whose text + # is not a classification we would want to key on. The FLAG is the classification, + # and it is the ONLY one: `_deadline` sets it before the cancel lands, so an error + # raised while it is unset is the database's own — however late it arrives. + if expired.is_set(): + raise _ResourceLimit("the statement outlived its per-statement budget") + raise + # Checked after the watchdog is disarmed, so the flag is final. A cancel can also land + # between the two calls above, or just as the second returns, and leave nothing to raise: + # the budget still elapsed, so the outcome is still a refusal rather than a result gathered + # past it. + if expired.is_set(): + raise _ResourceLimit("the statement outlived its per-statement budget") + except _ResourceLimit: + # Ahead of the catch-all below on purpose: our own marker must reach `execute_guarded` + # intact. Wrapped in an `ExecutorError` it would become a `failed`/`syntax` envelope, which + # is the database's outcome rather than the bound we imposed. + raise except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: @@ -1342,11 +1377,12 @@ def execute_guarded( **Every path returns exactly one ``Envelope``, and this function is TOTAL** — nothing is raised out of here for a caller to interpret, so the subprocess ``main`` and the in-process MCP handler - cannot disagree about what happened. The six outcomes: + cannot disagree about what happened. The seven outcomes: * the read-only gate refuses -> ``refused`` carrying that gate's ``Refusal`` * ``_model_safety`` returns a Refusal -> ``refused`` carrying it verbatim * ``_model_safety`` returns an int -> ``refused`` carrying the interim ``model_safety`` + * the per-statement deadline fired -> ``refused`` carrying ``resource_limit`` * ``executor.execute`` raises -> ``failed`` carrying a classified ``Failure`` * anything else raises -> ``failed``/``other``, generic message, raw to the log * the statement ran -> ``ok`` carrying the ``ExecResult`` @@ -1400,6 +1436,28 @@ def execute_guarded( # does not accept) fails the present-iff check in `Envelope.__post_init__`, and that is a # broken adapter, not a reason for the chokepoint to raise at its caller. return _envelope("ok", data=result) + except _ResourceLimit: + # AHEAD of both handlers below: `_ResourceLimit` is an `ExecutorError` sibling, not a + # subclass, but the catch-all would swallow it and report a bound WE imposed as an + # unclassified server break. This is the per-statement bound the contract reserves + # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a + # fix we can honestly name, unlike the supervisor's kill of a child that never returned. + # + # The budget is re-resolved rather than carried on the marker: the marker stays a plain + # exception, and nothing between the engine call and here can change the env var or the + # request-scoped ContextVar the resolver reads, so it is the same number the watchdog used. + timeout_s = _resolve_timeout_s() + return _envelope("refused", refusal=refuse( + RULE_RESOURCE_LIMIT, + # The configured number belongs here — it is a deployment setting, not a data value, and + # a bound the caller cannot see is one it cannot plan around. The remediation names only + # what would make THIS statement executable: on the served path the caller is an + # assistant with no shell and no deployment, so naming the environment variable would be + # advice it cannot take, addressed to someone who is not reading. + detail=f"The statement ran longer than the {timeout_s}s limit and was cancelled.", + remediation="Narrow the time range, reduce the grouping, or add a selective filter, " + "then run it again.", + )) except ExecutorError as exc: # The classified branch: `msg` is authored by this module (a missing driver, a connect # failure, the credential-resolution remediation naming DATASOURCE_URL) or relayed from the diff --git a/plugins/agami/lib/guardrail.py b/plugins/agami/lib/guardrail.py index 0e7e437..8c0bc07 100644 --- a/plugins/agami/lib/guardrail.py +++ b/plugins/agami/lib/guardrail.py @@ -50,11 +50,13 @@ RULE_COLUMN_SCOPE = "column_scope" RULE_SELECT_STAR = "select_star" RULE_MODEL_UNAVAILABLE = "model_unavailable" -# Declared and pinned below, with NO producer today. The contract reserves it for a **per-statement** -# timeout the guard imposes — a bound whose subject is the statement, so "narrow the query" is a fix -# we can honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this -# rule: it stops a child that never returned, without knowing what the child was doing when it -# stopped, so it is a `failed`/`timeout` (see `FailureKind` below and contract §3). +# Produced by the **per-statement** timeout the executor imposes: a watchdog cancels a statement that +# outlives the configured budget, and `execute_sql.execute_guarded` turns that into this refusal. Its +# subject is the statement, which is what earns it a rule at all — "narrow the query" is a fix we can +# honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this rule: it +# stops a child that never returned, without knowing what the child was doing when it stopped, so it +# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Wired into SQLite first; the +# other engines follow, and until they do a statement on them is bounded only by that supervisor. RULE_RESOURCE_LIMIT = "resource_limit" RULE_UNPARSEABLE = "unparseable" @@ -90,8 +92,8 @@ RULE_SELECT_STAR: "out_of_scope", RULE_MODEL_UNAVAILABLE: "undetermined", # A bound we imposed, not a property of the statement: neither unsafe nor out of scope — we - # simply did not determine the answer within the bound. Pinned here while unproduced so the gate - # that eventually imposes a per-statement timeout fills a constant rather than inventing one. + # simply did not determine the answer within the bound. Pinned before it had a producer, so the + # gate that imposes the per-statement timeout filled a constant rather than inventing one. RULE_RESOURCE_LIMIT: "undetermined", RULE_UNPARSEABLE: "undetermined", RULE_MODEL_SAFETY: "undetermined", @@ -165,9 +167,12 @@ def refuse(rule: str, *, detail: str, remediation: str) -> Refusal: in connect, credential resolution or model load, where "narrow the query" is the wrong fix. So an unresponsive executor is `failed` / `timeout`, and its message says only that we stopped waiting (guardrail contract §3). A **per-statement** timeout is the other case — its subject IS the -statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT` — and nothing imposes one yet, which is -why that rule is pinned with no producer. Driver-level connect/login timeouts fold into the connect -failure the executor already reports as `auth` (exit 4), because that is what the driver raises. +statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT`, and the executor now imposes one: a +watchdog cancels the statement through the driver and the outcome leaves the chokepoint on the +refusal channel rather than this one. The two therefore coexist and stay distinguishable in the audit +trail, which is the whole point of splitting them. Driver-level connect/login timeouts fold into the +connect failure the executor already reports as `auth` (exit 4), because that is what the driver +raises. `column_not_found`, `table_not_found`, `permission` and `network` are DECLARED BUT UNREACHABLE: producing them means parsing driver text, and sanitizing driver text belongs to the error-hardening diff --git a/tests/test_ace035_guardrail_audit.py b/tests/test_ace035_guardrail_audit.py index e62c9d4..d81d7ab 100644 --- a/tests/test_ace035_guardrail_audit.py +++ b/tests/test_ace035_guardrail_audit.py @@ -694,9 +694,10 @@ def test_killing_an_unresponsive_executor_is_a_failure_we_cannot_attribute(env, authority of a guardrail decision. `failed` / `timeout` claims exactly what we know — we stopped waiting — and the value-free message says only that. - `RULE_RESOURCE_LIMIT` stays declared and pinned in `REASON_FOR_RULE` for the per-statement bound - a later gate imposes, which IS a refusal because its subject is the statement. It has no producer - today, and this branch must not become one. + `RULE_RESOURCE_LIMIT` belongs to the per-statement bound the executor now imposes, which IS a + refusal because its subject is the statement. That the rule finally has a producer is precisely + why this branch is worth holding: the two bounds now coexist, and the one that cannot say what it + stopped must not drift into borrowing the vocabulary of the one that can. """ def _timed_out(cmd, **kwargs): raise subprocess.TimeoutExpired(cmd, kwargs.get("timeout", 240)) @@ -720,13 +721,15 @@ def _timed_out(cmd, **kwargs): assert row["datasource"] == PROFILE and row["question"] == QUESTION -def test_no_gate_produces_the_resource_limit_refusal(env, monkeypatch): - """`RULE_RESOURCE_LIMIT` is declared with NO producer, and the supervisor kill is where it would - creep back in — so drive that branch and assert the rule is absent from what the caller gets. +def test_the_supervisor_kill_does_not_borrow_the_resource_limit_refusal(env, monkeypatch): + """The supervisor kill is where `RULE_RESOURCE_LIMIT` would creep in — so drive that branch and + assert the rule is absent from what the caller gets, and from the row recording it. - The pin is on `REASON_FOR_RULE`, not on the absence: the rule stays in the contract's table so - the gate that eventually imposes a per-statement timeout fills a constant instead of inventing a - string. What must not happen is a branch borrowing it because the word "timeout" fits. + Renamed rather than retired: the rule DOES have a producer now (the executor's per-statement + deadline), so "no gate produces it" is no longer the claim. The claim that survives is the sharper + half and always was — a branch must not borrow this rule merely because the word "timeout" fits. + The pin is on `REASON_FOR_RULE`, not on the absence, which is why this test needed no weakening + when the producer landed: it asserts what this branch says, not what the contract contains. """ assert guardrail.RULE_RESOURCE_LIMIT in guardrail.REASON_FOR_RULE diff --git a/tests/test_ace035_no_enumeration.py b/tests/test_ace035_no_enumeration.py index 8bcb31c..2a124d8 100644 --- a/tests/test_ace035_no_enumeration.py +++ b/tests/test_ace035_no_enumeration.py @@ -15,18 +15,20 @@ schema-listing endpoint, reachable by anyone who can send one deliberately-wrong statement — which is the recon surface a later slice exists to close. -**This file is lock-in, not a fix.** All five gates are echo-only today; that was verified when they -were converted, and it is asserted here so a later reword cannot quietly turn a refusal into a -listing. Unlike the rest of this spec's tests, the property this one pins is expected to hold on -`446cc20` (the pre-conversion base) as well — and it does. It was checked by running this file's -model, canaries, vectors and scanner against a materialized `446cc20` tree, driving all five rules -through all four routes below and scanning the 20 refusal bodies the base produced: all clean. (The -file cannot be *collected* at that commit — `guardrail` does not exist there and the tool edge -speaks `{"error": {kind, remediation}}` rather than the Envelope — so the base run drops only the -two shape assertions, `status == "refused"` and `refusal.rule == …`, and keeps the scanner -verbatim.) A future failure here is therefore a live disclosure bug, not a conversion regression. - -**What this file covers.** Every field of the serialized tool-edge body, for the five refusal rules +**This file is lock-in, not a fix.** All five gates were echo-only when they were converted, and the +per-statement timeout that has since joined them is too; it is asserted here so a later reword cannot +quietly turn a refusal into a listing. Unlike the rest of this spec's tests, the property this one +pins is expected to hold on `446cc20` (the pre-conversion base) as well — and it does. It was checked +by running this file's model, canaries, vectors and scanner against a materialized `446cc20` tree, +driving the five rules that existed then through all four routes below and scanning the 20 refusal +bodies the base produced: all clean. (The `resource_limit` vector came later and has no counterpart +there — nothing imposed a per-statement bound at that commit. The file cannot be *collected* at that +commit either — `guardrail` does not exist there and the tool edge speaks +`{"error": {kind, remediation}}` rather than the Envelope — so the base run drops only the two shape +assertions, `status == "refused"` and `refusal.rule == …`, and keeps the scanner verbatim.) A future +failure here is therefore a live disclosure bug, not a conversion regression. + +**What this file covers.** Every field of the serialized tool-edge body, for the six refusal rules and for a `failed` — across both surfaces and both execution paths. `failure.message` is in scope because it is part of that body: it was the one field the scanner never saw, and it is where a PostgreSQL `HINT: Perhaps you meant to reference the column "orders.internal_ref".` arrives, which @@ -178,6 +180,11 @@ def declared(tmp_path, monkeypatch): monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(artifacts)) monkeypatch.setenv("DATASOURCE_URL__ACME", f"sqlite:///{warehouse}") + # The smallest per-statement budget the resolver accepts, for the `resource_limit` vector: it is + # the one statement here that reaches the executor and has to be stopped there. Every other + # statement in this file is refused before execution or rejected instantly by the database, so a + # short budget costs them nothing. + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") # Local, not hosted: the disk model is the one the gates use. (`model_unavailable` needs the # hosted signal and gets its own fixture.) monkeypatch.delenv("AGAMI_DB_URL", raising=False) @@ -297,18 +304,34 @@ def _route_http(sql: str, profile: str = PROFILE) -> dict: # The sentinel # --------------------------------------------------------------------------- +# A statement that passes every gate and then runs for minutes: a recursive CTE with two billion +# iterations, filtered on a subquery over the declared table so the vector is the shape a real +# runaway query has rather than a synthetic spin. `orders` and `id` are therefore the caller's own, +# and echoing them back is legitimate; the canaries stay unsent, which is what makes a refusal that +# reached for the model a failing assertion here. +_RUNAWAY_SQL = ( + "WITH RECURSIVE burn(n) AS (" + "SELECT 1 UNION ALL SELECT n + 1 FROM burn WHERE n < 2000000000" + ") SELECT count(n) AS c FROM burn WHERE n > (SELECT count(id) FROM orders)" +) + # One statement per rule, chosen so the rule under test is the FIRST gate to fire: -# read_only — a denied keyword, refused at the tool edge before a profile is resolved -# table_scope — an undeclared table (runs before the star ban and the column gate) -# select_star — a declared table, projected with `*` -# column_scope — a declared table, an undeclared column +# read_only — a denied keyword, refused at the tool edge before a profile is resolved +# table_scope — an undeclared table (runs before the star ban and the column gate) +# select_star — a declared table, projected with `*` +# column_scope — a declared table, an undeclared column +# resource_limit — the odd one out: every gate PASSES it, and the fixture's one-second budget stops +# it inside the executor. So this row scans the one refusal that is produced after +# the model has been loaded and consulted — the state in which a "helpful" message +# has the declared surface closest to hand. # `audit_trail` and `ref_no` are the caller's own inventions, so echoing them back is legitimate; -# nothing the MODEL declares appears in any of them except `orders`. +# nothing the MODEL declares appears in any of them except `orders` and `id`. VECTORS = ( (guardrail.RULE_READ_ONLY, "DELETE FROM orders"), (guardrail.RULE_TABLE_SCOPE, "SELECT ref_no FROM audit_trail"), (guardrail.RULE_SELECT_STAR, "SELECT * FROM orders"), (guardrail.RULE_COLUMN_SCOPE, "SELECT ref_no FROM orders"), + (guardrail.RULE_RESOURCE_LIMIT, _RUNAWAY_SQL), ) UNAVAILABLE_SQL = "SELECT id FROM orders" @@ -348,7 +371,7 @@ def _assert_echo_only(body: dict, sql: str) -> None: @pytest.mark.parametrize(("rule", "sql", "route"), _MATRIX, ids=_MATRIX_IDS) def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal(declared, rule, sql, route): - """Five rules — four here, `model_unavailable` below — across both surfaces and both execution + """Six rules — five here, `model_unavailable` below — across both surfaces and both execution paths. The fork column is not redundant with in-process: the child serializes the refusal to stderr and @@ -706,13 +729,6 @@ def test_echoing_the_callers_own_identifier_is_allowed(): "than refusing, so nothing constructs this rule yet; turning that fail-open into a refusal " "is the unparseable-statement slice's job, and the refusal it introduces needs a vector here." ), - guardrail.RULE_RESOURCE_LIMIT: ( - "No producer. It is reserved for the per-statement timeout a later gate imposes, and the " - "one branch that might borrow it — the subprocess supervisor killing an unresponsive " - "executor — is a `failed`/`timeout` instead, pinned by " - "test_ace035_guardrail_audit.py::test_no_gate_produces_the_resource_limit_refusal. A rule " - "nothing emits cannot leak; the day something emits it, it needs a vector here." - ), guardrail.RULE_MODEL_SAFETY: ( "Reachable, but its detail is authored as static prose at a single construction site in " "`execute_sql.execute_guarded` — the unconverted `_model_safety` branches hand back a bare " @@ -744,7 +760,7 @@ def test_every_rule_and_every_route_is_covered(): assert not covered & set(_NO_VECTOR) assert all(reason.strip() for reason in _NO_VECTOR.values()) assert set(ROUTES) == {"in_process", "fork", "stdio", "http"} - assert len(_MATRIX) == len(VECTORS) * len(ROUTES) == 16 + assert len(_MATRIX) == len(VECTORS) * len(ROUTES) == 20 # The `failed` channel is covered by its own matrix rather than this one, because it has no # rule. Its vector must not be one of these: a statement that a gate refuses would report on the # refusal channel a second time and leave `failure.message` unscanned again, which is exactly diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index 1312685..99881a4 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -1,19 +1,37 @@ -"""Per-statement timeout — config resolution and the deadline primitive. +"""Per-statement timeout — config resolution, the deadline primitive, and the refusal it produces. `_resolve_timeout_s` answers "how long may one statement run", from `AGAMI_SQL_TIMEOUT_S`, the `_timeout_override` ContextVar, and a 30s default; unlike the row cap it complains out loud when the env value is present but unparseable. `_deadline` is the watchdog those seconds feed: it fires an Event and calls a cancel callable when a block outlives its budget, and disarms cleanly when it does -not. Nothing in the engines calls either yet. +not. + +The second half of this file proves the whole contract end to end on ONE engine — SQLite, chosen +because it is in-process, needs no network, and `sqlite3.Connection.interrupt()` is a genuine cancel +rather than a polite request. A statement that outlives its budget is cancelled, unwinds on the +internal `_ResourceLimit` marker, and leaves `execute_guarded` as a `refused` Envelope carrying +`resource_limit` — with no partial data, a detail that quotes the configured budget, and a +remediation addressed to whoever can actually act on it. + +**The classification is the FLAG, and only the flag.** A cancelled SQLite statement raises +`OperationalError("interrupted")`, so neither the error text nor the elapsed clock can be the test: +both are properties an ordinary database error can have by coincidence, and reading either one would +mean an unlucky query gets told to narrow itself when nothing timed out. `_deadline` sets its Event +*before* the cancel lands, so "did WE stop this?" is answerable without inference — and it is +asserted here in both directions. """ from __future__ import annotations +import contextlib +import json import logging +import sqlite3 import sys import threading import time from pathlib import Path +from types import SimpleNamespace import pytest @@ -23,6 +41,7 @@ sys.path.insert(0, str(PKG_SRC)) import execute_sql # noqa: E402 +import guardrail # noqa: E402 # Short enough that the whole file stays sub-second, long enough that a loaded CI runner still # schedules the timer thread before the assertion runs. @@ -218,3 +237,408 @@ def test_the_resource_limit_marker_is_a_plain_exception(): assert issubclass(execute_sql._ResourceLimit, Exception) with pytest.raises(execute_sql._ResourceLimit): raise execute_sql._ResourceLimit("statement exceeded its budget") + + +# -------------------------------------------------------------------------------------------- +# The deadline, wired into SQLite — the refusal, end to end +# -------------------------------------------------------------------------------------------- + +PROFILE = "analytics" + +# The smallest budget the resolver accepts, since it deals in whole seconds. Every end-to-end test +# below therefore costs about a second of wall clock, which is the price of driving a real cancel +# through a real driver rather than asserting on a stub. +_BUDGET_S = 1 + +# A recursive CTE with no physical table and no termination in reach: two billion iterations, which +# on this machine is minutes of pure CPU. Sized that way on purpose — "bounded rather than hanging" +# is only proved if the unbounded run would take far longer than the assertion allows. +# +# It also has to reach the executor, so it is written to pass every gate ahead of it: it opens with +# WITH…SELECT (read-only), names no physical table (`burn` is a CTE, which the table-scope gate +# excludes by construction), projects no star, and its one column binds to a CTE rather than to a +# declared table. +_RUNAWAY_SQL = ( + "WITH RECURSIVE burn(n) AS (" + "SELECT 1 UNION ALL SELECT n + 1 FROM burn WHERE n < 2000000000" + ") SELECT count(n) AS c FROM burn" +) + + +@pytest.fixture +def warehouse(tmp_path, monkeypatch): + """A real SQLite warehouse reachable as profile `analytics`, and nothing else configured. + + `no_safety=True` on the direct calls below skips the semantic-model pass, so this fixture only + has to satisfy credential resolution. The audit test further down needs the fuller install and + builds it itself. + """ + path = tmp_path / "warehouse.db" + con = sqlite3.connect(path) + con.execute("CREATE TABLE orders (id INTEGER)") + con.commit() + con.close() + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", f"sqlite:///{path}") + monkeypatch.delenv("AGAMI_DB_URL", raising=False) + monkeypatch.delenv("APP_DATABASE_URL", raising=False) + return path + + +def _guarded(sql: str) -> object: + """`execute_guarded` over the built-in executor — the single chokepoint, driven directly.""" + return execute_sql.execute_guarded( + sql, PROFILE, None, executor=execute_sql.BUILTIN_EXECUTOR, no_safety=True + ) + + +def test_a_runaway_statement_is_cancelled_rather_than_left_to_run(warehouse): + """The headline: a statement that would run for minutes is stopped at its budget and comes back + as a refusal naming the rule the contract reserves for a bound we imposed. + + The elapsed assertion is the one that would still fail if the deadline were never armed — without + it a test that merely waited out the query would look identical and pass in several minutes. + """ + execute_sql._timeout_override.set(_BUDGET_S) + + started = time.monotonic() + env = _guarded(_RUNAWAY_SQL) + elapsed = time.monotonic() - started + + assert elapsed < 20, f"the statement ran {elapsed:.1f}s against a {_BUDGET_S}s budget" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + # Neither unsafe nor out of scope: we simply did not determine the answer within the bound. + assert env.refusal.reason == "undetermined" + assert env.refusal.reason == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + + +def test_a_cancelled_statement_yields_no_partial_data(warehouse): + """Whatever rows the engine had gathered when the watchdog fired are not an answer. + + A truncated result presented as a result is the failure mode the bounded-fetch work already + guards against on the row axis; on the time axis the answer is stronger — there is no data at + all, and `Envelope.__post_init__` enforces that a refusal cannot carry any. + """ + execute_sql._timeout_override.set(_BUDGET_S) + + env = _guarded(_RUNAWAY_SQL) + + assert env.status == "refused" + assert env.data is None + assert env.failure is None + + +class _ResourceLimitExecutor: + """A `ports.Executor` that raises the internal marker, standing in for an engine whose watchdog + fired. Used where the subject is the REFUSAL TEXT rather than the cancel, so the assertion does + not have to pay a second of wall clock to reach it.""" + + def execute(self, vetted_sql: str, creds: dict, *, profile: str): + raise execute_sql._ResourceLimit("the statement outlived its per-statement budget") + + +def test_the_detail_quotes_the_configured_budget(warehouse): + """A bound the caller cannot see is one it cannot plan around, so the number is in the detail. + + The configured value is not a data value: it is a deployment setting, and stating it discloses + nothing about the database or its contents. Asserted against a distinctive budget rather than the + default, so a hard-coded `30s` in the message cannot pass. + """ + execute_sql._timeout_override.set(7) + + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_ResourceLimitExecutor(), no_safety=True, + ) + + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert "7s" in env.refusal.detail, env.refusal.detail + + +def test_the_remediation_names_no_deployment_environment_variable(warehouse): + """The remediation has to be addressed to whoever is reading it. + + On the served path that is an assistant holding a statement, with no shell, no deployment and no + way to set an environment variable — so "raise AGAMI_SQL_TIMEOUT_S" is advice aimed past the + caller at an operator who is not in the conversation, and it reads as a fix while being + unfollowable. What is left has to be something that would make THIS statement executable. + """ + execute_sql._timeout_override.set(_BUDGET_S) + + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_ResourceLimitExecutor(), no_safety=True, + ) + + authored = f"{env.refusal.detail} {env.refusal.remediation}" + assert "AGAMI_SQL_TIMEOUT_S" not in authored, authored + assert "AGAMI_" not in authored, authored # no sibling deployment var either + assert "environ" not in authored and "env var" not in authored.lower(), authored + # And it is still actionable — it names something to change about the statement. + assert env.refusal.remediation.strip() + + +# -------------------------------------------------------------------------------------------- +# The clock covers the fetch, not just the execute +# -------------------------------------------------------------------------------------------- + + +class _SlowFetchCursor: + """A cursor whose `execute` returns at once and whose single `fetchmany` blocks until cancelled. + + That is the shape of a streaming / server-side cursor, where the scan happens on the PULL rather + than on the call: `_collect_cursor` issues one `fetchmany(cap + 1)`, and on such a cursor that + one call is the whole query. A deadline that stopped at `execute` would bound the cheap half. + """ + + def __init__(self, cancelled: threading.Event): + self.description = [("c",)] + self._cancelled = cancelled + self.executed: str | None = None + + def execute(self, sql: str) -> None: + self.executed = sql + + def fetchmany(self, n: int): + # Bounded so a deadline that never reaches the fetch fails this test rather than hanging it. + if not self._cancelled.wait(30): + raise AssertionError("the fetch was never cancelled — the clock stopped at execute") + raise sqlite3.OperationalError("interrupted") + + +class _FakeConnection: + """The driver surface `_run_sqlite` uses: a cursor, a real `interrupt`, and a close.""" + + def __init__(self, cursor_factory): + self.cancelled = threading.Event() + self.cursor_obj = cursor_factory(self.cancelled) + self.closed = False + + def cursor(self): + return self.cursor_obj + + def interrupt(self) -> None: + self.cancelled.set() + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def fake_sqlite(monkeypatch): + """Replace `sqlite3.connect` so a test can choose exactly where the time goes. + + A real slow query cannot separate execute from fetch — sqlite decides that — and driving one + would also make these tests as slow as the thing they measure. `_run_sqlite` does its own + `import sqlite3`, which resolves through `sys.modules`, so patching the module attribute is + enough to reach it. + """ + def _install(cursor_factory, *, connect_delay_s: float = 0.0): + conn = _FakeConnection(cursor_factory) + + def _connect(path, *a, **kw): + if connect_delay_s: + time.sleep(connect_delay_s) + return conn + + monkeypatch.setattr(sqlite3, "connect", _connect) + return conn + + return _install + + +def test_a_slow_fetch_is_bounded_even_when_the_execute_returned_at_once(warehouse, fake_sqlite): + """The clock covers the whole statement, fetch included — an explicit criterion, not a bonus. + + Bounding only `execute` would leave the common streaming shape unbounded: the driver returns + immediately and the engine scans while the caller pulls. The cancel has to land on the fetch, and + the refusal has to be the same one a slow execute produces. + """ + execute_sql._timeout_override.set(_BUDGET_S) + conn = fake_sqlite(_SlowFetchCursor) + + started = time.monotonic() + env = _guarded("SELECT c FROM orders") + elapsed = time.monotonic() - started + + assert conn.cursor_obj.executed == "SELECT c FROM orders" # the execute really did run, and fast + assert elapsed < 20, f"the fetch ran {elapsed:.1f}s against a {_BUDGET_S}s budget" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert conn.closed # the connection is still released on the refusing path + + +# -------------------------------------------------------------------------------------------- +# The classification is the flag alone — asserted in both directions +# -------------------------------------------------------------------------------------------- +# +# `interrupt()` makes the in-flight statement raise `OperationalError("interrupted")`. That text is +# therefore NOT evidence of a timeout — a database is free to raise it for its own reasons, and an +# error that merely arrives late is not one we caused. Both vectors below carry that exact text with +# the flag unset, so anything keying on the message or on the clock would turn them green as +# refusals; only the flag distinguishes them. + + +class _ImmediateErrorCursor: + """Raises the very error a cancel provokes, straight away and unprompted.""" + + def __init__(self, cancelled: threading.Event): + self.description = None + self._cancelled = cancelled + + def execute(self, sql: str) -> None: + raise sqlite3.OperationalError("interrupted") + + def fetchmany(self, n: int): # pragma: no cover - execute raises first + raise AssertionError("unreachable") + + +def test_a_database_error_with_the_flag_unset_is_a_failure_not_a_refusal(warehouse, fake_sqlite): + """Direction (a): the watchdog never fired, so this is the database's outcome, not ours. + + A generous budget means the flag stays clear while the identical error text arrives. It must + unwind as an `ExecutorError` and leave the chokepoint as `failed`/`syntax` — a refusal here would + tell a caller to narrow a statement that never ran long at all. + """ + execute_sql._timeout_override.set(300) + fake_sqlite(_ImmediateErrorCursor) + + env = _guarded("SELECT c FROM orders") + + assert env.status == "failed" + assert env.failure.kind == "syntax" + assert env.refusal is None + + +@contextlib.contextmanager +def _deadline_that_never_fires(cancel, timeout_s): + """A watchdog that does not fire, whatever the block does. + + Substituted so a block can genuinely outlive its budget with the flag clear — the one state a + real `_deadline` cannot be put into, and the one that separates "the clock ran out" from "we + cancelled it". + """ + yield threading.Event() + + +class _LateErrorCursor: + """Runs past the budget, then raises the error a cancel would have provoked.""" + + def __init__(self, cancelled: threading.Event): + self.description = None + self._cancelled = cancelled + + def execute(self, sql: str) -> None: + time.sleep(_BUDGET_S * 1.5) + raise sqlite3.OperationalError("interrupted") + + def fetchmany(self, n: int): # pragma: no cover - execute raises first + raise AssertionError("unreachable") + + +def test_an_error_that_merely_arrives_late_is_still_a_failure(warehouse, fake_sqlite, monkeypatch): + """Direction (b): elapsed time is not the classifier either. + + Everything a clock-based or text-based reading would need is present — the budget is exceeded and + the message is literally "interrupted" — and the flag is clear. The verdict must still be + `failed`, because we did not stop this statement; it stopped on its own, late. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_that_never_fires) + execute_sql._timeout_override.set(_BUDGET_S) + fake_sqlite(_LateErrorCursor) + + started = time.monotonic() + env = _guarded("SELECT c FROM orders") + elapsed = time.monotonic() - started + + assert elapsed > _BUDGET_S, "the statement did not actually outlive its budget" + assert env.status == "failed" + assert env.failure.kind == "syntax" + assert env.refusal is None + + +# -------------------------------------------------------------------------------------------- +# The refusal is recorded +# -------------------------------------------------------------------------------------------- + + +@pytest.fixture +def audited(tmp_path, monkeypatch): + """A complete single-datasource install: an app database to audit into, a semantic model on + disk, and the same real warehouse. + + Fuller than the `warehouse` fixture because this one goes through the tool edge rather than + straight to `execute_guarded`, so the model pass runs for real — and with an app database + configured the executor reads that as the hosted signal and fails closed without a model. + """ + pytest.importorskip("pydantic") + pytest.importorskip("sqlglot") + yaml = pytest.importorskip("yaml") + from store import Store + + app_db = "sqlite://" + str(tmp_path / "app.db") + store = Store.connect(app_db) + store.run_migrations() + store.close() + + root = tmp_path / "artifacts" / PROFILE + (root / "subject_areas" / "sales" / "tables").mkdir(parents=True) + (root / "datasource.yaml").write_text(yaml.safe_dump( + {"datasource": "Shop", "version": 1, "subject_areas": ["subject_areas/sales"]})) + (root / "subject_areas" / "sales" / "subject_area.yaml").write_text(yaml.safe_dump( + {"name": "sales", "tables": [ + {"storage_connection": "c", "schema": "public", "table": "orders"}]})) + (root / "subject_areas" / "sales" / "tables" / "orders.yaml").write_text(yaml.safe_dump({ + "name": "orders", "schema": "public", "storage_connection": "c", "grain": ["id"], + "description": "orders", + "columns": [{"name": "id", "type": "integer", "primary_key": True}], + })) + + path = tmp_path / "warehouse.db" + con = sqlite3.connect(path) + con.execute("CREATE TABLE orders (id INTEGER)") + con.commit() + con.close() + + monkeypatch.setenv("AGAMI_DB_URL", app_db) + monkeypatch.delenv("APP_DATABASE_URL", raising=False) + monkeypatch.delenv("AGAMI_ORG_ID", raising=False) + monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "artifacts")) + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", f"sqlite:///{path}") + return SimpleNamespace(app_db=app_db) + + +def test_the_refusal_is_written_to_the_audit_trail(audited): + """A decision we made against a caller's statement is exactly the row a reviewer comes looking + for, so this outcome is audited like every other — and keyed by the id the caller was handed. + + Driven through the tool edge with the built-in executor injected, which is the in-process path + the hosted server runs: one real cancel, one real serializer, one real sink. + """ + import tools + from store import Store + + tools.set_injected_executor(execute_sql.BUILTIN_EXECUTOR) + execute_sql._timeout_override.set(_BUDGET_S) + try: + body = json.loads(tools.tool_execute_sql({"sql": _RUNAWAY_SQL, "datasource": PROFILE, + "raw_query": "how many"})) + finally: + tools.set_injected_executor(None) + + assert body["status"] == "refused", body + assert body["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT, body + + store = Store.connect(audited.app_db) + try: + rows = store.query("SELECT id, status, reason, rule FROM query_executions") + finally: + store.close() + + assert len(rows) == 1, rows + (row,) = rows + assert row["id"] == body["audit_id"] # the answer and its record name the same id + assert row["status"] == "refused" + assert row["rule"] == guardrail.RULE_RESOURCE_LIMIT + assert row["reason"] == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] From 50404c12e045f3a196296fd1f82ef745e957b8db Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 15:52:03 -0700 Subject: [PATCH 03/10] feat(executor): the deadline on every engine, and the native bounds that back it (S3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit S2 proved the contract on one engine. The remaining nine now run under the same watchdog, and three of them also carry a native server-side bound behind it. **Each engine names its own cancel, explicitly.** A duck-typed probe was considered and is provably wrong, in both directions. `pymysql.Connection` has NO `cancel()` — its methods are `close`, `_force_close`, `kill` — so a probe falls through to `close()`, which sends COM_QUIT down the very socket the blocked statement owns and can itself block; `_force_close()` shuts the socket outright and is the one that unblocks. In the other direction `oracledb.Connection` DOES have `cancel()`, so a probe ordering `cancel` first would silently pick it for Oracle and never notice that Snowflake, Databricks and Trino put theirs on the CURSOR. So: postgres/redshift `conn.cancel`, mysql/mariadb `conn._force_close`, sqlserver the `_mssql` connection's `cancel` behind pymssql's private `_conn` (its DB-API connection has none), oracle `conn.cancel`, databricks and trino `cur.cancel`, duckdb `conn.interrupt` — where `cur` IS `conn`, because `conn.execute` hands back the connection itself. Snowflake departs from the plan on evidence: the connector exposes no `cancel()` on either the connection or the cursor. `SnowflakeCursor.abort_query(qid)` is the only public abort, so the cancel resolves the cursor's `sfqid` and calls it, and no-ops before the statement has a query id — the window the session parameter covers. Every engine wraps BOTH `cur.execute` and `_collect_cursor` in the deadline, resolves its budget ONCE per call, and converts to `_ResourceLimit` only when the Event is set. An `except _ResourceLimit: raise` now sits ahead of all ten `code=5` catch-alls, BigQuery included even though nothing there raises it today: the rule is that no engine may relabel the marker, and a rule checkable on nine of ten is not one. **Postgres needed a structural fix before the deadline could work there at all.** Closing a server-side cursor sends `CLOSE agami_bounded`. On the timeout path the transaction is already aborted, so that statement raises in turn — and from inside `__exit__`, where a new exception REPLACES the one being propagated. The marker vanished and the catch-all reported every Postgres timeout as a `failed` envelope. The named cursor comes out of its `with` and is closed by hand, with that raise swallowed; `with conn` stays, because it is the transaction and the rollback. **Native bounds on exactly three engines, at `timeout_s + 5`.** The skew is the design: our watchdog fires first, so its flag stays the sole classification signal, and the native bound is the backstop for a process that dies mid-statement. Postgres sets `SET LOCAL statement_timeout` in ms on the same transaction, on a regular cursor, before the named cursor declares — not the libpq `options` startup parameter, which a transaction-mode pooler can reject at connect. Snowflake sets `STATEMENT_TIMEOUT_IN_SECONDS` via `session_parameters`, where an abandoned statement still bills credits. BigQuery sets `job_timeout_ms` and gets NO watchdog: there is no connection to cancel and the blocking call is `job.result()`, reached only after `client.query()` returns, so at the instant a watchdog would fire there is nothing in hand. That residual is recorded in the code and asserted in the tests — on BigQuery a client-side stall yields `failed`, not `resource_limit`. No other engine gets one. Tests are table-driven over all ten `_run_*` functions, against fake driver modules installed in `sys.modules` (the engines do their own `import `), so the whole matrix runs on a machine with no database drivers at all: * every engine re-raises the marker; and, inversely, no engine mistakes an ordinary driver error for it. A coverage test asserts the table names every `_run_*` in the module, so an engine added without the clause fails the suite. * the named cancel is the one actually invoked. The fake connection exposes EVERY cancel-shaped method these drivers have — `cancel`, `interrupt`, `_force_close`, `kill`, `close`, a cursor `cancel` and a cursor `abort_query` — each logging a distinct marker, so reaching for the wrong one lands on the wrong marker. * Postgres: the marker survives a named-cursor close that raises (fails without the fix, as `failed`/`syntax` carrying "current transaction is aborted"), the transaction still rolls back, and `SET LOCAL` runs on the same transaction before the named cursor with `(budget + 5) * 1000` ms. * the three native bounds and their values; and that no fourth engine grew one. * a real cartesian bomb bounded on a real in-process DuckDB. It skips here, since the suite deliberately installs no DB drivers. `test_ace044_bounded_fetch.py`'s Postgres fake gains `cancel()` and a `params` argument on `execute` — the calls the engine now really makes. Neither assertion in it changed. Spec: ACE-038 Co-Authored-By: Claude Opus 5 (1M context) --- packages/agami-core/src/execute_sql.py | 242 ++++++++- plugins/agami/lib/execute_sql.py | 242 ++++++++- tests/test_ace038_timeout.py | 698 +++++++++++++++++++++++++ tests/test_ace044_bounded_fetch.py | 12 +- 4 files changed, 1141 insertions(+), 53 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 62c6b80..1f396ff 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -486,20 +486,62 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Postgres connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # Bound outside the try so the `finally` can close it whether or not the declare got that far. + cur = None try: + # `with conn` is the TRANSACTION, not the connection: leaving it by an exception rolls back, + # which is why it stays even though the cursor no longer lives inside a `with` of its own. with conn: + # `SET LOCAL statement_timeout` is the native server-side backstop. It is + # TRANSACTION-scoped, so it has to run on THIS transaction and before the named cursor + # declares — a regular client-side cursor, because a named cursor may only ever run the + # one query it was declared for. The libpq `options` startup parameter would be the + # other way to set it and is deliberately not used: a transaction-mode connection pooler + # can reject an unknown startup parameter, which would break the connect outright rather + # than bound the statement. + with conn.cursor() as bound_cur: + bound_cur.execute( + "SET LOCAL statement_timeout = %s", + ((timeout_s + _NATIVE_BOUND_SKEW_S) * 1000,), # the setting is in milliseconds + ) # A server-side (named) cursor so the row cap bounds TRANSFER, not just what we write: # psycopg2's default client-side cursor buffers the ENTIRE result before we can fetchmany, # so a runaway result would still be pulled whole. The named cursor streams from the # server in bounded batches (ACE-038). Read-only SELECTs (the only thing the guard admits) # are exactly what a server-side cursor supports. - with conn.cursor(name="agami_bounded") as cur: - cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window - cur.execute(sql) - result = _collect_cursor(cur) + cur = conn.cursor(name="agami_bounded") + cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window + # `connection.cancel()` is psycopg2's cancel: it opens a second connection and sends the + # libpq cancel request, so it is safe to call from the watchdog thread while this one is + # blocked in the driver. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Postgres execution error: {e}", code=5) finally: + # The named cursor is closed HERE, by hand, rather than by a `with` around it. Closing one + # sends `CLOSE agami_bounded` to the server, and when we are unwinding on the timeout marker + # the transaction is already aborted, so that statement raises in turn — from inside + # `__exit__`, where a raised exception REPLACES the one being propagated. The refusal would + # reach the chokepoint as an ordinary driver error and be reported as a failure. Swallowing + # it costs nothing: by this point the transaction has ended and the server-side portal is + # gone with it, so the close is a courtesy rather than the thing that frees the resource. + if cur is not None: + try: + cur.close() + except Exception: + pass conn.close() return result @@ -523,10 +565,26 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"MySQL connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: with conn.cursor() as cur: - cur.execute(sql) - result = _collect_cursor(cur) + # pymysql's connection has NO `cancel()`. Its two cancel-shaped methods are `close()`, + # which sends COM_QUIT down the very socket the blocked statement owns — and so can + # itself block on a connection that is already stuck — and `_force_close()`, which + # closes the socket outright. Only the second one actually unblocks a statement in + # flight, so it is named here explicitly rather than reached for by shape. + with _deadline(conn._force_close, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: @@ -551,11 +609,20 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: "Add one to /local/credentials.", code=2, ) + # Resolved before the connect, because the native backstop is a SESSION parameter and so has to + # be handed to the connect call itself. + timeout_s = _resolve_timeout_s() conn_kwargs: dict[str, Any] = { "account": creds["account"], "user": creds["user"], "client_session_keep_alive": False, "login_timeout": 15, + # Snowflake's native server-side bound, set behind our watchdog by the usual skew so the + # watchdog wins and the flag stays the only classification. On a warehouse the backstop is + # worth more than elsewhere: a statement nobody is waiting for still bills credits. + "session_parameters": { + "STATEMENT_TIMEOUT_IN_SECONDS": timeout_s + _NATIVE_BOUND_SKEW_S, + }, } for k in ("password", "warehouse", "database", "schema", "role", "authenticator"): if creds.get(k): @@ -566,8 +633,32 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"Snowflake connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + + def cancel_snowflake_statement() -> None: + """Ask Snowflake to abort the statement this cursor is running. + + Neither the connection nor the cursor has a `cancel()` — verified against the connector's + own source, where the only public abort is `SnowflakeCursor.abort_query(qid)` and the + connection's cancel helpers are private. The query id is what identifies the statement to + abort, and it exists only once the statement has been submitted; before that there is + nothing to abort and the session parameter above is what bounds the call. + """ + qid = cur.sfqid + if qid: + cur.abort_query(qid) + + with _deadline(cancel_snowflake_statement, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Snowflake execution error: {e}", code=5) finally: @@ -637,9 +728,20 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"BigQuery client init failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # `job_timeout_ms` is BigQuery's native server-side bound, and on this engine it is the ONLY + # bound: there is no watchdog cancel here, deliberately. BigQuery hands out no connection or + # cursor to cancel, and the call that blocks is `job.result()`, which is reached only AFTER + # `client.query()` has returned — so at the instant a watchdog would fire there is nothing in + # hand to stop. Stated plainly, and accepted rather than papered over: on BigQuery a client-side + # stall comes back as a `failed` envelope, not a `resource_limit` refusal. The query itself is + # still stopped by the bound below, which is what keeps a runaway from scanning on unattended. + # The usual skew is kept so this engine's number matches every other engine's. + job_config_kwargs: dict[str, Any] = { + "job_timeout_ms": (timeout_s + _NATIVE_BOUND_SKEW_S) * 1000, + } # If `dataset` was set, prefix unqualified table references via the # default_dataset job config so the SQL can omit `..` - job_config_kwargs: dict[str, Any] = {} if creds.get("dataset"): try: job_config_kwargs["default_dataset"] = f"{project}.{creds['dataset']}" @@ -648,15 +750,17 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: cap = _resolve_row_cap() try: - if job_config_kwargs: - job_config = bigquery.QueryJobConfig(**job_config_kwargs) - job = client.query(sql, job_config=job_config) - else: - job = client.query(sql) + job_config = bigquery.QueryJobConfig(**job_config_kwargs) + job = client.query(sql, job_config=job_config) # BigQuery has no DB-API cursor, so it can't funnel through `_collect_cursor`; apply the # same bounded-fetch cap here. `max_results=cap+1` bounds what the API returns (transfer), # and the (cap+1)th row flags truncation — the never-silent guarantee holds for BigQuery too. results = job.result(max_results=cap + 1) # waits for completion; raises on error + except _ResourceLimit: + # Nothing in this engine raises the marker today, and this clause is still not optional: the + # rule is that NO engine may relabel it as a driver error, and the clause below would. It is + # what makes the rule checkable across all ten engines rather than nine. + raise except Exception as e: raise ExecutorError(f"BigQuery execution error: {e}", code=5) @@ -707,14 +811,14 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: # and it is the ONLY one: `_deadline` sets it before the cancel lands, so an error # raised while it is unset is the database's own — however late it arrives. if expired.is_set(): - raise _ResourceLimit("the statement outlived its per-statement budget") + raise _ResourceLimit(_OUTLIVED_BUDGET) raise # Checked after the watchdog is disarmed, so the flag is final. A cancel can also land # between the two calls above, or just as the second returns, and leave nothing to raise: # the budget still elapsed, so the outcome is still a refusal rather than a result gathered # past it. if expired.is_set(): - raise _ResourceLimit("the statement outlived its per-statement budget") + raise _ResourceLimit(_OUTLIVED_BUDGET) except _ResourceLimit: # Ahead of the catch-all below on purpose: our own marker must reach `execute_guarded` # intact. Wrapped in an `ExecutorError` it would become a `failed`/`syntax` envelope, which @@ -742,10 +846,26 @@ def _run_sqlserver(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"SQL Server connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # pymssql's DB-API `Connection` has no `cancel()` — verified against pymssql 2.3, whose + # connection exposes only close/commit/cursor/rollback/bulk_copy, and whose cursor exposes + # none either. The cancel lives one layer down, on the `_mssql.MSSQLConnection` that + # connection wraps, reachable only as the private `_conn`; that object's `cancel()` sends the + # TDS attention packet, which is the thing that actually stops a statement in flight. + with _deadline(conn._conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"SQL Server execution error: {e}", code=5) finally: @@ -772,10 +892,23 @@ def _run_oracle(creds: dict[str, str], sql: str) -> ExecResult: conn = oracledb.connect(user=creds["user"], password=creds["password"], dsn=dsn) except Exception as e: raise ExecutorError(f"Oracle connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # `oracledb.Connection.cancel()` breaks out of the call in progress on that connection — + # verified against python-oracledb 3.x, where it is on the CONNECTION and not on the cursor. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Oracle execution error: {e}", code=5) finally: @@ -803,10 +936,24 @@ def _run_databricks(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Databricks connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The cancel is on the CURSOR here, not the connection — verified against + # databricks-sql-connector 4.x, whose `Cursor.cancel()` posts a cancel for the operation that + # cursor is running while its connection has none. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Databricks execution error: {e}", code=5) finally: @@ -835,10 +982,23 @@ def _run_trino(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Trino connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # Trino's cancel is on the CURSOR — verified against the trino client, whose + # `Cursor.cancel()` sends the coordinator a DELETE for the query that cursor started. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Trino execution error: {e}", code=5) finally: @@ -860,9 +1020,25 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: conn = duckdb.connect(path, read_only=True) except Exception as e: raise ExecutorError(f"DuckDB open failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: - cur = conn.execute(sql) - result = _collect_cursor(cur) + # `conn.interrupt()` is DuckDB's cancel, and it is on the connection — which is just as well, + # because on this engine there is no cursor to reach for until the execute has returned: + # `conn.execute` HANDS BACK THE CONNECTION ITSELF, so `cur` below IS `conn`. The deadline + # therefore has to be armed around the execute as well as the fetch, and on an in-process + # engine the execute is where the scan happens. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur = conn.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"DuckDB execution error: {e}", code=5) finally: @@ -898,6 +1074,14 @@ def _resolve_row_cap() -> int: _DEFAULT_TIMEOUT_S = 30 # wall-clock seconds one statement may run before the watchdog cancels it +# How far BEHIND our own watchdog a NATIVE server-side bound is set, on the three engines that have +# one. The skew is the whole point: our watchdog fires at the budget, the engine's own bound five +# seconds later, so the watchdog always wins the race and its Event stays the SOLE classification +# signal. Reverse the order and a server-side kill would beat us to the statement, the flag would be +# clear, and the refusal would come back as an ordinary database failure. The native bound is a +# backstop for the one case the watchdog cannot cover — our process dying with a statement in +# flight, which would otherwise leave the engine scanning for nobody. +_NATIVE_BOUND_SKEW_S = 5 # Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the # HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the # context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one @@ -940,6 +1124,12 @@ class _ResourceLimit(Exception): boundary.""" +# The marker's message, single-sourced because every engine raises it. It is diagnostic text, not +# caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own +# detail, so nothing a caller reads comes from here. +_OUTLIVED_BUDGET = "the statement outlived its per-statement budget" + + @contextlib.contextmanager def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threading.Event]: """Arm a watchdog that calls `cancel` if the wrapped block outlives `timeout_s`, and yield the diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 62c6b80..1f396ff 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -486,20 +486,62 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Postgres connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # Bound outside the try so the `finally` can close it whether or not the declare got that far. + cur = None try: + # `with conn` is the TRANSACTION, not the connection: leaving it by an exception rolls back, + # which is why it stays even though the cursor no longer lives inside a `with` of its own. with conn: + # `SET LOCAL statement_timeout` is the native server-side backstop. It is + # TRANSACTION-scoped, so it has to run on THIS transaction and before the named cursor + # declares — a regular client-side cursor, because a named cursor may only ever run the + # one query it was declared for. The libpq `options` startup parameter would be the + # other way to set it and is deliberately not used: a transaction-mode connection pooler + # can reject an unknown startup parameter, which would break the connect outright rather + # than bound the statement. + with conn.cursor() as bound_cur: + bound_cur.execute( + "SET LOCAL statement_timeout = %s", + ((timeout_s + _NATIVE_BOUND_SKEW_S) * 1000,), # the setting is in milliseconds + ) # A server-side (named) cursor so the row cap bounds TRANSFER, not just what we write: # psycopg2's default client-side cursor buffers the ENTIRE result before we can fetchmany, # so a runaway result would still be pulled whole. The named cursor streams from the # server in bounded batches (ACE-038). Read-only SELECTs (the only thing the guard admits) # are exactly what a server-side cursor supports. - with conn.cursor(name="agami_bounded") as cur: - cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window - cur.execute(sql) - result = _collect_cursor(cur) + cur = conn.cursor(name="agami_bounded") + cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window + # `connection.cancel()` is psycopg2's cancel: it opens a second connection and sends the + # libpq cancel request, so it is safe to call from the watchdog thread while this one is + # blocked in the driver. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Postgres execution error: {e}", code=5) finally: + # The named cursor is closed HERE, by hand, rather than by a `with` around it. Closing one + # sends `CLOSE agami_bounded` to the server, and when we are unwinding on the timeout marker + # the transaction is already aborted, so that statement raises in turn — from inside + # `__exit__`, where a raised exception REPLACES the one being propagated. The refusal would + # reach the chokepoint as an ordinary driver error and be reported as a failure. Swallowing + # it costs nothing: by this point the transaction has ended and the server-side portal is + # gone with it, so the close is a courtesy rather than the thing that frees the resource. + if cur is not None: + try: + cur.close() + except Exception: + pass conn.close() return result @@ -523,10 +565,26 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"MySQL connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: with conn.cursor() as cur: - cur.execute(sql) - result = _collect_cursor(cur) + # pymysql's connection has NO `cancel()`. Its two cancel-shaped methods are `close()`, + # which sends COM_QUIT down the very socket the blocked statement owns — and so can + # itself block on a connection that is already stuck — and `_force_close()`, which + # closes the socket outright. Only the second one actually unblocks a statement in + # flight, so it is named here explicitly rather than reached for by shape. + with _deadline(conn._force_close, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: @@ -551,11 +609,20 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: "Add one to /local/credentials.", code=2, ) + # Resolved before the connect, because the native backstop is a SESSION parameter and so has to + # be handed to the connect call itself. + timeout_s = _resolve_timeout_s() conn_kwargs: dict[str, Any] = { "account": creds["account"], "user": creds["user"], "client_session_keep_alive": False, "login_timeout": 15, + # Snowflake's native server-side bound, set behind our watchdog by the usual skew so the + # watchdog wins and the flag stays the only classification. On a warehouse the backstop is + # worth more than elsewhere: a statement nobody is waiting for still bills credits. + "session_parameters": { + "STATEMENT_TIMEOUT_IN_SECONDS": timeout_s + _NATIVE_BOUND_SKEW_S, + }, } for k in ("password", "warehouse", "database", "schema", "role", "authenticator"): if creds.get(k): @@ -566,8 +633,32 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"Snowflake connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + + def cancel_snowflake_statement() -> None: + """Ask Snowflake to abort the statement this cursor is running. + + Neither the connection nor the cursor has a `cancel()` — verified against the connector's + own source, where the only public abort is `SnowflakeCursor.abort_query(qid)` and the + connection's cancel helpers are private. The query id is what identifies the statement to + abort, and it exists only once the statement has been submitted; before that there is + nothing to abort and the session parameter above is what bounds the call. + """ + qid = cur.sfqid + if qid: + cur.abort_query(qid) + + with _deadline(cancel_snowflake_statement, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Snowflake execution error: {e}", code=5) finally: @@ -637,9 +728,20 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"BigQuery client init failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # `job_timeout_ms` is BigQuery's native server-side bound, and on this engine it is the ONLY + # bound: there is no watchdog cancel here, deliberately. BigQuery hands out no connection or + # cursor to cancel, and the call that blocks is `job.result()`, which is reached only AFTER + # `client.query()` has returned — so at the instant a watchdog would fire there is nothing in + # hand to stop. Stated plainly, and accepted rather than papered over: on BigQuery a client-side + # stall comes back as a `failed` envelope, not a `resource_limit` refusal. The query itself is + # still stopped by the bound below, which is what keeps a runaway from scanning on unattended. + # The usual skew is kept so this engine's number matches every other engine's. + job_config_kwargs: dict[str, Any] = { + "job_timeout_ms": (timeout_s + _NATIVE_BOUND_SKEW_S) * 1000, + } # If `dataset` was set, prefix unqualified table references via the # default_dataset job config so the SQL can omit `..` - job_config_kwargs: dict[str, Any] = {} if creds.get("dataset"): try: job_config_kwargs["default_dataset"] = f"{project}.{creds['dataset']}" @@ -648,15 +750,17 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: cap = _resolve_row_cap() try: - if job_config_kwargs: - job_config = bigquery.QueryJobConfig(**job_config_kwargs) - job = client.query(sql, job_config=job_config) - else: - job = client.query(sql) + job_config = bigquery.QueryJobConfig(**job_config_kwargs) + job = client.query(sql, job_config=job_config) # BigQuery has no DB-API cursor, so it can't funnel through `_collect_cursor`; apply the # same bounded-fetch cap here. `max_results=cap+1` bounds what the API returns (transfer), # and the (cap+1)th row flags truncation — the never-silent guarantee holds for BigQuery too. results = job.result(max_results=cap + 1) # waits for completion; raises on error + except _ResourceLimit: + # Nothing in this engine raises the marker today, and this clause is still not optional: the + # rule is that NO engine may relabel it as a driver error, and the clause below would. It is + # what makes the rule checkable across all ten engines rather than nine. + raise except Exception as e: raise ExecutorError(f"BigQuery execution error: {e}", code=5) @@ -707,14 +811,14 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: # and it is the ONLY one: `_deadline` sets it before the cancel lands, so an error # raised while it is unset is the database's own — however late it arrives. if expired.is_set(): - raise _ResourceLimit("the statement outlived its per-statement budget") + raise _ResourceLimit(_OUTLIVED_BUDGET) raise # Checked after the watchdog is disarmed, so the flag is final. A cancel can also land # between the two calls above, or just as the second returns, and leave nothing to raise: # the budget still elapsed, so the outcome is still a refusal rather than a result gathered # past it. if expired.is_set(): - raise _ResourceLimit("the statement outlived its per-statement budget") + raise _ResourceLimit(_OUTLIVED_BUDGET) except _ResourceLimit: # Ahead of the catch-all below on purpose: our own marker must reach `execute_guarded` # intact. Wrapped in an `ExecutorError` it would become a `failed`/`syntax` envelope, which @@ -742,10 +846,26 @@ def _run_sqlserver(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"SQL Server connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # pymssql's DB-API `Connection` has no `cancel()` — verified against pymssql 2.3, whose + # connection exposes only close/commit/cursor/rollback/bulk_copy, and whose cursor exposes + # none either. The cancel lives one layer down, on the `_mssql.MSSQLConnection` that + # connection wraps, reachable only as the private `_conn`; that object's `cancel()` sends the + # TDS attention packet, which is the thing that actually stops a statement in flight. + with _deadline(conn._conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"SQL Server execution error: {e}", code=5) finally: @@ -772,10 +892,23 @@ def _run_oracle(creds: dict[str, str], sql: str) -> ExecResult: conn = oracledb.connect(user=creds["user"], password=creds["password"], dsn=dsn) except Exception as e: raise ExecutorError(f"Oracle connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # `oracledb.Connection.cancel()` breaks out of the call in progress on that connection — + # verified against python-oracledb 3.x, where it is on the CONNECTION and not on the cursor. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Oracle execution error: {e}", code=5) finally: @@ -803,10 +936,24 @@ def _run_databricks(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Databricks connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The cancel is on the CURSOR here, not the connection — verified against + # databricks-sql-connector 4.x, whose `Cursor.cancel()` posts a cancel for the operation that + # cursor is running while its connection has none. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Databricks execution error: {e}", code=5) finally: @@ -835,10 +982,23 @@ def _run_trino(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Trino connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # Trino's cancel is on the CURSOR — verified against the trino client, whose + # `Cursor.cancel()` sends the coordinator a DELETE for the query that cursor started. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Trino execution error: {e}", code=5) finally: @@ -860,9 +1020,25 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: conn = duckdb.connect(path, read_only=True) except Exception as e: raise ExecutorError(f"DuckDB open failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: - cur = conn.execute(sql) - result = _collect_cursor(cur) + # `conn.interrupt()` is DuckDB's cancel, and it is on the connection — which is just as well, + # because on this engine there is no cursor to reach for until the execute has returned: + # `conn.execute` HANDS BACK THE CONNECTION ITSELF, so `cur` below IS `conn`. The deadline + # therefore has to be armed around the execute as well as the fetch, and on an in-process + # engine the execute is where the scan happens. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur = conn.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"DuckDB execution error: {e}", code=5) finally: @@ -898,6 +1074,14 @@ def _resolve_row_cap() -> int: _DEFAULT_TIMEOUT_S = 30 # wall-clock seconds one statement may run before the watchdog cancels it +# How far BEHIND our own watchdog a NATIVE server-side bound is set, on the three engines that have +# one. The skew is the whole point: our watchdog fires at the budget, the engine's own bound five +# seconds later, so the watchdog always wins the race and its Event stays the SOLE classification +# signal. Reverse the order and a server-side kill would beat us to the statement, the flag would be +# clear, and the refusal would come back as an ordinary database failure. The native bound is a +# backstop for the one case the watchdog cannot cover — our process dying with a statement in +# flight, which would otherwise leave the engine scanning for nobody. +_NATIVE_BOUND_SKEW_S = 5 # Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the # HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the # context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one @@ -940,6 +1124,12 @@ class _ResourceLimit(Exception): boundary.""" +# The marker's message, single-sourced because every engine raises it. It is diagnostic text, not +# caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own +# detail, so nothing a caller reads comes from here. +_OUTLIVED_BUDGET = "the statement outlived its per-statement budget" + + @contextlib.contextmanager def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threading.Event]: """Arm a watchdog that calls `cancel` if the wrapped block outlives `timeout_s`, and yield the diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index 99881a4..53f41b7 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -30,6 +30,7 @@ import sys import threading import time +import types from pathlib import Path from types import SimpleNamespace @@ -642,3 +643,700 @@ def test_the_refusal_is_written_to_the_audit_trail(audited): assert row["status"] == "refused" assert row["rule"] == guardrail.RULE_RESOURCE_LIMIT assert row["reason"] == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + + +# -------------------------------------------------------------------------------------------- +# Every engine, under one table +# -------------------------------------------------------------------------------------------- +# +# S2 proved the contract on SQLite. The other nine now run under the same deadline, and each names +# its OWN cancel — because there is no method a duck-typed probe could look for that is right +# everywhere. pymysql's connection has no `cancel()` at all and a probe would fall through to +# `close()`, which sends COM_QUIT down the socket the blocked statement owns; oracledb's connection +# DOES have `cancel()`, so a probe that tried `cancel` first would silently pick it for Oracle and +# never notice that Snowflake and Databricks put theirs on the cursor. The fakes below stand in for +# drivers this environment does not install, so the whole matrix runs on every machine. + + +def _module(name: str, **attrs: object): + """A stand-in driver module. The engine functions do their own `import `, which resolves + through `sys.modules`, so an entry there is enough to reach them.""" + mod = types.ModuleType(name) + for key, value in attrs.items(): + setattr(mod, key, value) + return mod + + +class _RecordingCursor: + """A DB-API cursor that logs every call, so a test can assert WHICH method a cancel reached.""" + + def __init__(self, log: list, *, name: str | None = None, on_execute=None, + close_raises: bool = False): + self.description = [("c",)] + self.name = name + self.itersize = 0 + self.sfqid = "01fake-0000-0000-0000-000000000000" # Snowflake sets a query id at submission + self._log = log + self._on_execute = on_execute + self._close_raises = close_raises + + def execute(self, sql: str, params=None) -> None: + self._log.append(("execute", sql, params)) + # Only the caller's statement is armed to fail. Postgres runs `SET LOCAL statement_timeout` + # on its own cursor first, and that one has to succeed — it runs before anything has gone + # wrong, and a test about the timeout path must not accidentally break the setup for it. + if self._on_execute is not None and not sql.startswith("SET "): + raise self._on_execute() + + def fetchmany(self, n: int): + self._log.append(("fetchmany", n)) + return [(1,)] + + def cancel(self) -> None: + self._log.append(("cursor.cancel",)) + + def abort_query(self, qid: str) -> bool: + self._log.append(("cursor.abort_query", qid)) + return True + + def close(self) -> None: + self._log.append(("cursor.close", self.name)) + if self._close_raises: + raise RuntimeError("current transaction is aborted, commands ignored") + + def __enter__(self): + return self + + def __exit__(self, *exc_info) -> bool: + self.close() + return False + + +class _RecordingConnection: + """Every cancel-shaped method a driver in this module might expose, each logging a DISTINCT + marker — so asserting the log pins down exactly which one the engine chose.""" + + def __init__(self, *, on_execute=None, named_close_raises: bool = False): + self.log: list = [] + self.connect_kwargs: dict = {} + self.job_config_kwargs: dict = {} + self._on_execute = on_execute + self._named_close_raises = named_close_raises + self.cursors: list[_RecordingCursor] = [] + self.executed_cursor = None + # pymssql's DB-API connection delegates its cancel to the `_mssql` connection it wraps. + self._conn = SimpleNamespace(cancel=lambda: self.log.append(("mssql.cancel",))) + + def cursor(self, name: str | None = None, **kwargs): + self.log.append(("cursor", name)) + cur = _RecordingCursor( + self.log, + name=name, + on_execute=self._on_execute, + close_raises=self._named_close_raises and name is not None, + ) + self.cursors.append(cur) + return cur + + # The per-driver cancels, each distinguishable in the log. + def cancel(self) -> None: + self.log.append(("conn.cancel",)) + + def interrupt(self) -> None: + self.log.append(("conn.interrupt",)) + + def _force_close(self) -> None: + self.log.append(("conn._force_close",)) + + def kill(self, thread_id: int) -> None: # pragma: no cover - present so a probe could find it + self.log.append(("conn.kill",)) + + def close(self) -> None: + self.log.append(("conn.close",)) + + # DuckDB's `execute` hands back the CONNECTION, so the connection is also a cursor there. + def execute(self, sql: str, params=None): + self.log.append(("execute", sql, params)) + self.executed_cursor = self + if self._on_execute is not None: + raise self._on_execute() + return self + + @property + def description(self): + return [("c",)] + + def fetchmany(self, n: int): + self.log.append(("fetchmany", n)) + return [(1,)] + + # `with conn` is the TRANSACTION on the Postgres path. + def __enter__(self): + self.log.append(("txn.enter",)) + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + self.log.append(("txn.exit", exc_type is not None)) + return False + + +class _FakeBigQueryResults: + def __init__(self): + self.schema = [SimpleNamespace(name="c")] + + def __iter__(self): + return iter([(1,)]) + + +def _install_simple(module_name: str, connect_attr: str = "connect"): + """Installer for a driver reached as `import ` + `.connect(...)`.""" + def _install(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute, named_close_raises=named_close_raises) + + def _connect(*args, **kwargs): + conn.connect_kwargs = kwargs + return conn + + monkeypatch.setitem(sys.modules, module_name, _module(module_name, **{connect_attr: _connect})) + return conn + + return _install + + +def _install_sqlite(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + monkeypatch.setattr(sqlite3, "connect", lambda *a, **kw: conn) + return conn + + +def _install_snowflake(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + def _connect(**kwargs): + conn.connect_kwargs = kwargs + return conn + + connector = _module("snowflake.connector", connect=_connect) + # `import snowflake.connector` short-circuits on `sys.modules` without the machinery ever setting + # the attribute, so the parent module has to carry it explicitly. + monkeypatch.setitem(sys.modules, "snowflake", _module("snowflake", connector=connector)) + monkeypatch.setitem(sys.modules, "snowflake.connector", connector) + return conn + + +def _install_databricks(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + def _connect(**kwargs): + conn.connect_kwargs = kwargs + return conn + + dbsql = _module("databricks.sql", connect=_connect) + monkeypatch.setitem(sys.modules, "databricks", _module("databricks", sql=dbsql)) + monkeypatch.setitem(sys.modules, "databricks.sql", dbsql) + return conn + + +def _install_trino(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + def _connect(**kwargs): + conn.connect_kwargs = kwargs + return conn + + monkeypatch.setitem( + sys.modules, "trino", _module("trino", dbapi=SimpleNamespace(connect=_connect)), + ) + return conn + + +def _install_bigquery(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + class _Job: + def result(self, max_results=None): + conn.log.append(("job.result", max_results)) + return _FakeBigQueryResults() + + class _Client: + def __init__(self, **kwargs): + conn.connect_kwargs = kwargs + + def query(self, sql, job_config=None): + conn.log.append(("execute", sql, None)) + if on_execute is not None: + raise on_execute() + return _Job() + + def _job_config(**kwargs): + conn.job_config_kwargs = kwargs + return SimpleNamespace(**kwargs) + + bigquery = _module("google.cloud.bigquery", Client=_Client, QueryJobConfig=_job_config) + oauth2 = _module("google.oauth2", service_account=_module("google.oauth2.service_account")) + monkeypatch.setitem(sys.modules, "google", _module("google")) + monkeypatch.setitem(sys.modules, "google.cloud", _module("google.cloud", bigquery=bigquery)) + monkeypatch.setitem(sys.modules, "google.cloud.bigquery", bigquery) + monkeypatch.setitem(sys.modules, "google.oauth2", oauth2) + monkeypatch.setitem(sys.modules, "google.oauth2.service_account", oauth2.service_account) + return conn + + +class _EngineCase: + """One engine's fake driver, its credentials, and the cancel it is required to reach.""" + + def __init__(self, fn, install, creds: dict, cancel: str | None): + self.fn = fn + self.install = install + self.creds = creds + self.cancel = cancel # the log marker the watchdog's cancel must produce; None = no watchdog + + def run(self, sql: str = "SELECT c FROM orders"): + return self.fn(self.creds, sql) + + +_SQL_CREDS = {"host": "db.example", "port": "5432", "user": "u", "password": "p", + "database": "shop"} + +ENGINE_CASES = { + "postgres": _EngineCase( + execute_sql._run_postgres, _install_simple("psycopg2"), dict(_SQL_CREDS), "conn.cancel"), + "mysql": _EngineCase( + execute_sql._run_mysql, _install_simple("pymysql"), dict(_SQL_CREDS), "conn._force_close"), + "snowflake": _EngineCase( + execute_sql._run_snowflake, _install_snowflake, + {"account": "acct", "user": "u", "password": "p"}, "cursor.abort_query"), + "bigquery": _EngineCase( + execute_sql._run_bigquery, _install_bigquery, {"project": "proj"}, None), + "sqlite": _EngineCase( + execute_sql._run_sqlite, _install_sqlite, {"path": "warehouse.db"}, "conn.interrupt"), + "sqlserver": _EngineCase( + execute_sql._run_sqlserver, _install_simple("pymssql"), + {"host": "db.example", "user": "u", "password": "p"}, "mssql.cancel"), + "oracle": _EngineCase( + execute_sql._run_oracle, _install_simple("oracledb"), + {"user": "u", "password": "p", "dsn": "db.example/shop"}, "conn.cancel"), + "databricks": _EngineCase( + execute_sql._run_databricks, _install_databricks, + {"host": "db.example", "http_path": "/sql/1", "token": "t"}, "cursor.cancel"), + "trino": _EngineCase( + execute_sql._run_trino, _install_trino, + {"host": "db.example", "user": "u"}, "cursor.cancel"), + "duckdb": _EngineCase( + execute_sql._run_duckdb, _install_simple("duckdb"), {"path": ":memory:"}, "conn.interrupt"), +} + + +def test_the_engine_table_covers_every_engine_this_module_has(): + """The guard that makes the table below a rule rather than a list. A new engine added without an + `except _ResourceLimit: raise` would otherwise ship silently reporting its refusals as driver + errors — this fails the moment a `_run_*` exists that the matrix does not name.""" + in_module = {n for n in dir(execute_sql) if n.startswith("_run_")} + assert in_module == {case.fn.__name__ for case in ENGINE_CASES.values()} + assert len(ENGINE_CASES) == 10 + + +def _raise_marker(): + return execute_sql._ResourceLimit(execute_sql._OUTLIVED_BUDGET) + + +@pytest.mark.parametrize("engine", sorted(ENGINE_CASES)) +def test_every_engine_re_raises_the_marker_instead_of_relabelling_it(engine, monkeypatch): + """The highest-value assertion in the slice, and the one a copy-paste omission fails. + + Each engine ends in `except Exception as e: raise ExecutorError(..., code=5)`. Without an + `except _ResourceLimit: raise` ahead of it, that catch-all swallows our own marker and the + chokepoint reports a bound WE imposed as the database's failure — a `failed` envelope telling the + caller their SQL is broken, when in fact it merely ran long. + """ + case = ENGINE_CASES[engine] + case.install(monkeypatch, on_execute=_raise_marker) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + +@pytest.mark.parametrize("engine", sorted(ENGINE_CASES)) +def test_no_engine_mistakes_an_ordinary_driver_error_for_the_marker(engine, monkeypatch): + """The other direction, on the same matrix: with the watchdog never fired, a driver error is the + database's outcome and has to stay one. An engine that raised the marker here would tell every + caller with a typo to narrow their query.""" + case = ENGINE_CASES[engine] + case.install(monkeypatch, on_execute=lambda: RuntimeError("relation does not exist")) + + with pytest.raises(execute_sql.ExecutorError) as exc: + case.run() + assert exc.value.code == 5 + + +@contextlib.contextmanager +def _deadline_already_fired(cancel, timeout_s): + """A watchdog that has already fired by the time the block runs — the state a real one reaches + only after a genuinely slow statement, reproduced here without paying for one.""" + fired = threading.Event() + fired.set() + yield fired + + +_WATCHDOG_ENGINES = sorted(e for e, c in ENGINE_CASES.items() if c.cancel is not None) + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_a_driver_error_under_a_fired_watchdog_becomes_the_marker(engine, monkeypatch): + """The conversion itself, on all nine engines that have a watchdog: the flag is set, so whatever + the driver raised is the wreckage of OUR cancel and unwinds as the marker.""" + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES[engine] + case.install(monkeypatch, on_execute=lambda: RuntimeError("connection reset by peer")) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_a_cancel_that_lands_without_raising_is_still_a_refusal(engine, monkeypatch): + """A cancel can land between the execute and the fetch, or just as the fetch returns, and leave + nothing to raise. The budget still elapsed, so the outcome is still a refusal rather than a + result gathered past it — which is why every engine re-reads the flag after the block.""" + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES[engine] + case.install(monkeypatch) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + +# -------------------------------------------------------------------------------------------- +# The named cancel is the one that actually runs +# -------------------------------------------------------------------------------------------- + + +class _CancelRecorder: + """Stands in for `_deadline` and keeps the callable it was handed, so a test can invoke exactly + what the watchdog would have invoked and see where it lands.""" + + def __init__(self, log: list | None = None): + self.cancels: list = [] + self.budgets: list = [] + # When handed the driver's own log, the arming and disarming are recorded IN it, so a test + # can assert where the deadline sits relative to the calls it is supposed to bound. + self._log = log + + @contextlib.contextmanager + def __call__(self, cancel, timeout_s): + self.cancels.append(cancel) + self.budgets.append(timeout_s) + if self._log is not None: + self._log.append(("deadline.arm",)) + try: + yield threading.Event() + finally: + if self._log is not None: + self._log.append(("deadline.disarm",)) + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_each_engine_arms_its_own_named_cancel(engine, monkeypatch): + """Not "a cancel ran" but "THE cancel ran". The fake connection exposes every cancel-shaped + method any of these drivers has — `cancel`, `interrupt`, `_force_close`, `kill`, `close`, a + cursor `cancel` and a cursor `abort_query` — each logging a distinct marker, so an engine that + reached for the wrong one lands on the wrong marker and fails here.""" + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert len(recorder.cancels) == 1, "exactly one deadline per call, resolved once" + before = len(conn.log) + recorder.cancels[0]() + landed = [entry[0] for entry in conn.log[before:]] + assert landed == [case.cancel], f"{engine} cancelled via {landed}, expected {[case.cancel]}" + + +def test_the_mysql_cancel_forces_the_socket_shut_and_never_sends_quit(monkeypatch): + """Called out on its own because it is the case a duck-typed probe gets wrong and no test would + notice. `pymysql.Connection` has no `cancel()`, so a probe falls through to `close()` — which + writes COM_QUIT to the very socket the blocked statement owns and can block on it. Only + `_force_close()` shuts the socket outright, which is what unblocks the statement.""" + case = ENGINE_CASES["mysql"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + before = len(conn.log) + recorder.cancels[0]() + + during_cancel = [entry[0] for entry in conn.log[before:]] + assert during_cancel == ["conn._force_close"] + assert "conn.close" not in during_cancel + assert "conn.kill" not in during_cancel + + +def test_the_postgres_cancel_is_the_connections_own(monkeypatch): + """psycopg2's `connection.cancel()` opens a second connection and sends the libpq cancel request, + so it is safe to call while this thread is blocked inside the driver.""" + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + before = len(conn.log) + recorder.cancels[0]() + + assert [entry[0] for entry in conn.log[before:]] == ["conn.cancel"] + + +def test_the_snowflake_cancel_does_nothing_before_a_query_id_exists(monkeypatch): + """The abort is addressed to a query id, and Snowflake only issues one at submission. A cancel + firing in the sliver before that has nothing to abort, and must not raise on the way to finding + out — the session parameter is what bounds the statement in that window.""" + case = ENGINE_CASES["snowflake"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + conn.cursors[0].sfqid = None # rewind to "submitted nothing yet" + before = len(conn.log) + recorder.cancels[0]() + + assert conn.log[before:] == [] + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_each_engine_arms_the_deadline_with_the_resolved_budget(engine, monkeypatch): + """One resolution per call, so the budget the watchdog enforces and the number the refusal quotes + cannot be two different values.""" + execute_sql._timeout_override.set(11) + case = ENGINE_CASES[engine] + case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert recorder.budgets == [11] + + +def test_duckdb_cancels_through_the_connection_even_though_the_cursor_is_the_connection(monkeypatch): + """DuckDB's `execute` hands back the CONNECTION rather than a cursor, so `cur` and `conn` are the + same object and a cursor-side cancel would be indistinguishable from a connection-side one by + inspection. The cancel that works is `interrupt()`, and it has to be armed around the execute: + on an in-process engine the execute is where the scan happens, not the fetch.""" + case = ENGINE_CASES["duckdb"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder(conn.log) + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert conn.executed_cursor is conn # the premise: execute returned the connection itself + # The deadline is armed BEFORE the execute and covers the fetch as well — on an in-process + # engine the execute is where the scan happens, so arming it after would bound nothing. + kinds = [entry[0] for entry in conn.log] + assert kinds[:1] == ["deadline.arm"] + assert kinds.index("execute") < kinds.index("fetchmany") < kinds.index("deadline.disarm") + before = len(conn.log) + recorder.cancels[0]() + assert [entry[0] for entry in conn.log[before:]] == ["conn.interrupt"] + + +# -------------------------------------------------------------------------------------------- +# Postgres: the named cursor must not eat the marker +# -------------------------------------------------------------------------------------------- + + +class _PostgresExecutor: + """Runs the real `_run_postgres` against the fake driver, so the assertion is about the Envelope + the chokepoint produces rather than about the exception the engine raises.""" + + def __init__(self, creds: dict): + self._creds = creds + + def execute(self, vetted_sql: str, creds: dict, *, profile: str): + return execute_sql._run_postgres(self._creds, vetted_sql) + + +def test_the_postgres_marker_survives_a_named_cursor_close_that_raises(warehouse, monkeypatch): + """The structural bug this slice had to fix before the deadline could work on Postgres at all. + + Closing a server-side cursor sends `CLOSE agami_bounded`. On the timeout path the transaction is + already aborted, so that statement raises in turn — and with the cursor inside a `with`, it + raises from `__exit__`, where a new exception REPLACES the one being propagated. The marker + vanishes, the engine's own catch-all wraps the replacement in an `ExecutorError`, and Postgres + alone reports every timeout as a failure. The fix is to close it by hand and swallow that. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES["postgres"] + conn = case.install( + monkeypatch, + on_execute=lambda: RuntimeError("canceling statement due to user request"), + named_close_raises=True, + ) + + env = execute_sql.execute_guarded( + "SELECT c FROM orders", PROFILE, None, + executor=_PostgresExecutor(case.creds), no_safety=True, + ) + + assert [e for e in conn.log if e[0] == "cursor.close"], "the named cursor was never closed" + assert env.status == "refused", getattr(env, "failure", None) + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert conn.log[-1] == ("conn.close",) # and the connection is still released + + +def test_the_postgres_transaction_still_rolls_back_when_the_deadline_fires(monkeypatch): + """`with conn` stayed for a reason: it is the TRANSACTION, and leaving it by an exception rolls + back. Taking the cursor out of its own `with` must not take that with it.""" + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch, on_execute=lambda: RuntimeError("canceling statement")) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + assert ("txn.enter",) in conn.log + assert ("txn.exit", True) in conn.log # exited WITH an exception in flight, i.e. rolled back + + +# -------------------------------------------------------------------------------------------- +# The native server-side bounds — three engines, each one skew behind the watchdog +# -------------------------------------------------------------------------------------------- + +_BUDGET_FOR_NATIVE = 11 # distinctive, so a hard-coded 30 or 35 cannot pass + + +def test_the_skew_puts_the_native_bound_behind_the_watchdog(): + """The direction of the skew is the whole design. Ahead of the watchdog, a server-side kill would + win the race, the flag would be clear, and the refusal would arrive as a database failure.""" + assert execute_sql._NATIVE_BOUND_SKEW_S > 0 + + +def test_postgres_sets_the_native_bound_on_the_same_transaction_before_the_named_cursor(monkeypatch): + """`SET LOCAL` is transaction-scoped, so it is worthless unless it runs on the transaction the + statement will run in, and before it. Ordering is the assertion; the value is the other half.""" + execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch) + + case.run() + + kinds = [entry[0] for entry in conn.log] + set_local = next( + i for i, e in enumerate(conn.log) + if e[0] == "execute" and e[1].startswith("SET LOCAL statement_timeout") + ) + declared = conn.log.index(("cursor", "agami_bounded")) + statement = next( + i for i, e in enumerate(conn.log) if e[0] == "execute" and e[1] == "SELECT c FROM orders" + ) + + assert kinds.index("txn.enter") < set_local < declared < statement + # Same transaction: nothing committed or rolled back between the setting and the statement. + assert "txn.exit" not in kinds[:statement] + # The setting is in milliseconds, and it sits one skew behind our own budget. + assert conn.log[set_local][2] == ((_BUDGET_FOR_NATIVE + execute_sql._NATIVE_BOUND_SKEW_S) * 1000,) + + +def test_postgres_does_not_set_the_native_bound_through_the_connect_options(monkeypatch): + """The libpq `options` startup parameter is the other way to set `statement_timeout`, and it is + the wrong one here: a transaction-mode connection pooler can reject an unknown startup parameter, + which breaks the connect outright rather than bounding the statement.""" + execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch) + + case.run() + + assert "options" not in conn.connect_kwargs + + +def test_snowflake_sets_its_statement_timeout_as_a_session_parameter(monkeypatch): + execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + case = ENGINE_CASES["snowflake"] + conn = case.install(monkeypatch) + + case.run() + + session = conn.connect_kwargs["session_parameters"] + assert session["STATEMENT_TIMEOUT_IN_SECONDS"] == ( + _BUDGET_FOR_NATIVE + execute_sql._NATIVE_BOUND_SKEW_S + ) + + +def test_bigquery_sets_a_job_timeout_and_is_the_one_engine_with_no_cancel(monkeypatch): + """BigQuery's native bound is the ONLY bound it has, and that is a recorded residual rather than + an oversight: there is no connection to cancel, and the call that blocks is `job.result()`, which + is reached only after `client.query()` returns — so at the instant a watchdog would fire there is + nothing in hand to stop. A client-side stall here comes back `failed`, not `resource_limit`.""" + execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + case = ENGINE_CASES["bigquery"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert conn.job_config_kwargs["job_timeout_ms"] == ( + _BUDGET_FOR_NATIVE + execute_sql._NATIVE_BOUND_SKEW_S + ) * 1000 + assert recorder.cancels == [], "BigQuery arms no watchdog — see the residual above" + + +def test_bigquery_keeps_the_default_dataset_alongside_the_job_timeout(monkeypatch): + """The job config used to be built only when a default dataset was configured. Now it is always + built, and the pre-existing setting has to survive that.""" + case = ENGINE_CASES["bigquery"] + conn = case.install(monkeypatch) + + execute_sql._run_bigquery({"project": "proj", "dataset": "shop"}, "SELECT c FROM orders") + + assert conn.job_config_kwargs["default_dataset"] == "proj.shop" + assert "job_timeout_ms" in conn.job_config_kwargs + + +@pytest.mark.parametrize("engine", sorted(set(ENGINE_CASES) - {"postgres", "snowflake", "bigquery"})) +def test_no_other_engine_grows_a_native_bound(engine, monkeypatch): + """Exactly three engines get one. The rest are bounded by the watchdog alone, and inventing a + per-engine session setting for them would be a second, unasserted timeout to keep in step.""" + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch) + + case.run() + + statements = " ".join(e[1] for e in conn.log if e[0] == "execute" and isinstance(e[1], str)) + assert "timeout" not in statements.lower() + assert "session_parameters" not in conn.connect_kwargs + + +# -------------------------------------------------------------------------------------------- +# A real in-process bomb, on a real DuckDB +# -------------------------------------------------------------------------------------------- + + +def test_a_cartesian_bomb_is_bounded_on_a_real_duckdb(tmp_path): + """The fakes above prove the wiring; this proves the cancel. DuckDB is in-process like SQLite, so + a genuine `interrupt()` can be driven end to end with no network and no fixture warehouse — and a + cross join of two ten-million-row ranges is 10^14 rows, far beyond what the budget allows, so a + run that returns in time can only have been stopped.""" + duckdb = pytest.importorskip("duckdb", reason="duckdb is not installed in this environment") + + path = tmp_path / "warehouse.duckdb" + con = duckdb.connect(str(path)) + con.execute("CREATE TABLE orders (id INTEGER)") + con.close() + + execute_sql._timeout_override.set(_BUDGET_S) + started = time.monotonic() + with pytest.raises(execute_sql._ResourceLimit): + execute_sql._run_duckdb( + {"path": str(path)}, + "SELECT count(*) AS c FROM range(10000000) a, range(10000000) b", + ) + elapsed = time.monotonic() - started + + assert elapsed < 20, f"the statement ran {elapsed:.1f}s against a {_BUDGET_S}s budget" diff --git a/tests/test_ace044_bounded_fetch.py b/tests/test_ace044_bounded_fetch.py index 2c1d8e7..e4025a8 100644 --- a/tests/test_ace044_bounded_fetch.py +++ b/tests/test_ace044_bounded_fetch.py @@ -175,13 +175,19 @@ def __enter__(self): def __exit__(self, *a): return False - def execute(self, sql): + def execute(self, sql, params=None): + # `params` accepts the native server-side bound the engine now sets first + # (`SET LOCAL statement_timeout = %s`, on its own client-side cursor). The caller's + # statement runs last, so `seen["sql"]` still ends up holding it. seen["sql"] = sql def fetchmany(self, n): seen["fetchmany"] = n return self._rows[:n] + def close(self): + seen["closed"] = True + class FakeConn: def cursor(self, name=None): return FakeCur(name) @@ -192,6 +198,10 @@ def __enter__(self): def __exit__(self, *a): return False + def cancel(self): + # The per-statement watchdog arms `connection.cancel` on this engine; never fired here. + seen["cancelled"] = True + def close(self): pass From 3b585fc427ffa473559170e8835ef5c1acbc0f8f Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 16:08:19 -0700 Subject: [PATCH 04/10] feat(executor): bound every executor at the chokepoint, and derive the supervisor (S4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The inner deadline lives inside the BUILT-IN executor's ten engine functions, which leaves two holes the criterion "the limit applies identically on stdio and HTTP, and no executor escapes it" does not survive. `execute_guarded` now bounds `executor.execute` itself, and the subprocess supervisor stops being a number of its own. **An injected executor was entirely unbounded.** `tools` routes to `_run_in_process` with the injected executor whenever a consumer supplies one, which is what the hosted connection-reuse path does — so that deployment had no per-statement bound at all. And on the built-in path BigQuery has no watchdog (no connection object to cancel, and the blocking call is reached only after `client.query()` returns), so a client-side stall there was unbounded too. Both are closed by the SAME layer, deliberately: applying the outer bound only to non-built-in executors would leave the BigQuery hole open while looking closed. **The mechanism is a daemon worker joined with the budget, because this layer holds nothing it can cancel.** An arbitrary `Executor` exposes one method and no connection, so the only thing this code owns is its own WAIT. On expiry the worker is abandoned, not stopped: it is still inside the driver call and may hold a database connection — and a statement still running on the server — until that call returns on its own. That cost is stated in the code rather than papered over, and it is exactly why this bound is the outer one and the watchdog is set to win. The refusal says so too: it does not claim the statement was cancelled, because nothing was. **Exception type survives the thread hop**, which is the subtle part. `BaseException` is captured in the worker and re-raised in the caller, so `except _ResourceLimit` and `except ExecutorError` still match and a driver's deep `sys.exit` still reaches the `SystemExit` net in `tools._run_in_process` instead of vanishing into a silent wait. The call runs inside `copy_context()`: a new thread starts with an EMPTY context, so without it the request-scoped `_timeout_override` and `_max_rows_override` would read as unset and every in-process call would quietly fall back to the deployment defaults. **One refusal, however many layers are armed.** The outer marker is a `_ResourceLimit` subclass, so the single existing handler mints exactly one envelope and no second one is possible; when the inner watchdog fired, its marker is what arrives there and the outer join had already returned, long before its own budget. **The supervisor is now `budget + 60`, from the same resolver.** A hardcoded 240s became the FIRST bound to fire for any statement budget approaching it, turning a refusal that can name the statement into a `failed`/`timeout` that names nothing. Only the DURATION moved: an unresponsive child is still `failed`/`timeout` and never `resource_limit`, because all we observe is that the child stopped responding — it may have hung in connect, in credential resolution or in loading the model. The child inherits `os.environ` and `subprocess.run` passes no `env=`, so it resolves the identical budget. The four bounds are one ordered family from one resolved budget — watchdog < native (+5) < outer (+10) < supervisor (+60) — and the tests assert that as a SINGLE fact rather than one skew at a time, since the order is the property and four separate assertions would each stay green while the family drifted apart. Ten new tests: a blocking injected executor bounded and refused (verified to fail without the outer layer); the outer refusal's wording, which must not say "cancelled"; the inner refusal winning a real cancel race with both layers armed at their real values; `ExecutorError` and a plain `Exception` still landing in their own handlers with the same envelope and kind as before; both ContextVars visible on a thread asserted not to be the caller's; the supervisor's computed value at a RAISED budget (300 -> 360, where the fixed 240 would have fired a minute early) read off the `subprocess.run` call rather than waited out, and its verdict unchanged; and both surfaces — a real `python -m mcp_harness` subprocess and `TestClient` over `/mcp` — returning identical refusals for the same runaway statement. No existing test changed a claim. Two fixtures gained setup only: `_reset_override` now also resets the row cap (the worker copies the whole context, so a leaked cap would be visible to the next test's executor), and `audited` sets the two vars the HTTP surface needs to mint a token. A new autouse fixture clears `_INJECTED_EXECUTOR` around every test in the file, since `create_app()` installs one as a side effect. Spec: ACE-038 Co-Authored-By: Claude Opus 5 (1M context) --- packages/agami-core/src/execute_sql.py | 130 +++++++- packages/agami-core/src/tools.py | 13 +- plugins/agami/lib/execute_sql.py | 130 +++++++- tests/test_ace038_timeout.py | 406 ++++++++++++++++++++++++- 4 files changed, 657 insertions(+), 22 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 1f396ff..d53ef2d 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -63,7 +63,7 @@ import urllib.parse import uuid from collections.abc import Callable, Iterator -from contextvars import ContextVar +from contextvars import ContextVar, copy_context from dataclasses import asdict, dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -1082,6 +1082,26 @@ def _resolve_row_cap() -> int: # backstop for the one case the watchdog cannot cover — our process dying with a statement in # flight, which would otherwise leave the engine scanning for nobody. _NATIVE_BOUND_SKEW_S = 5 +# How far behind the watchdog the OUTER bound sits — the one `execute_guarded` puts around +# `executor.execute` itself, so an INJECTED executor is bounded too. The four bounds are one ordered +# family resolved from the same budget, innermost first: +# +# watchdog (timeout_s) < native (+5) < outer (+10) < supervisor (+60) +# +# and the order is the contract. The inner watchdog must always win, because it is the only layer +# that can attribute the stop to the statement and hand the caller the precise refusal; every layer +# behind it is a backstop for a failure the layer in front of it cannot see. Collapse the order and +# a bound we imposed comes back as something else — a native server-side kill reads as an ordinary +# database error, and a supervisor kill reads as `failed`/`timeout` naming nothing the caller can +# act on. Both are strictly worse answers to the same event. +_OUTER_BOUND_SKEW_S = 10 +# The OUTERMOST bound: how far behind the watchdog the subprocess supervisor in `tools` stops waiting +# for a forked child. It lives here, next to its three siblings, so the whole family is derived from +# one resolver and one place — the supervisor was a hardcoded 240s, which for any statement budget +# approaching it became the FIRST bound to fire and inverted the order. Its slack is large because it +# bounds a whole process (interpreter start, model load, credential resolution, connect) rather than +# a statement. +_SUPERVISOR_SKEW_S = 60 # Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the # HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the # context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one @@ -1124,10 +1144,24 @@ class _ResourceLimit(Exception): boundary.""" +class _OuterBoundExpired(_ResourceLimit): + """Raised when the OUTER bound expired: the executor never returned and we stopped waiting. + + A subclass, so the one `except _ResourceLimit` handler in `execute_guarded` produces exactly one + refusal for either layer and no second one can be minted. It stays distinguishable because the + two events are not the same event: the watchdog CANCELLED a statement it holds a connection to, + while this layer only abandoned its wait — so the sentence the caller reads differs, and saying + "cancelled" here would be false. + """ + + # The marker's message, single-sourced because every engine raises it. It is diagnostic text, not # caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own # detail, so nothing a caller reads comes from here. _OUTLIVED_BUDGET = "the statement outlived its per-statement budget" +# The outer marker's message, and diagnostic in the same way — it names the executor rather than the +# statement, because at this layer the statement is not the thing we observed. +_OUTLIVED_OUTER_BOUND = "the executor outlived the outer bound around it" @contextlib.contextmanager @@ -1548,6 +1582,62 @@ def _envelope( ) +def _execute_bounded( + executor: Executor, sql: str, creds: dict[str, str], *, profile: str +) -> ExecResult: + """Run ``executor.execute`` under the OUTER bound and return what it returned. + + This is the layer that makes "no executor escapes the limit" true. The inner deadline lives + inside the BUILT-IN executor's engine functions, so before this an INJECTED executor — the + hosted connection-reuse path — ran with no per-statement bound at all, and even on the built-in + path BigQuery has no watchdog (it has no connection object to cancel), so a client-side stall + there was unbounded too. Both are bounded here, and deliberately by the same layer: applying + this only to non-built-in executors would leave the BigQuery hole open while looking closed. + + **The mechanism is a worker thread, because this layer holds nothing it can cancel.** An + arbitrary ``Executor`` exposes one method and no connection, so the only thing this code owns is + its own WAIT — it starts the call on a daemon thread and joins with the budget. + + **The cost is a leaked worker, and it is real.** On expiry the thread is still inside the + driver call, and it stays there: nothing here cancels it, and it may hold a database connection + (and, on the server side, a running statement) until that call returns on its own. It is a + daemon so it cannot hold the interpreter open at exit, and that is the whole of the mitigation. + This is a bound on how long a CALLER waits, not a promise that the work stopped — which is + exactly why the inner watchdog, the layer that really can cancel, is set to fire first. + + Exceptions cross the thread boundary with their ORIGINAL type, re-raised here. That is what + keeps the handlers in ``execute_guarded`` correct: ``_ResourceLimit`` still reaches the refusal + branch and ``ExecutorError`` still reaches the classified one, rather than every executor + failure arriving as a worker that finished having produced nothing. ``BaseException`` is caught + rather than ``Exception`` for the same reason: a driver's deep ``sys.exit`` raises ``SystemExit``, + which ``tools._run_in_process`` nets one layer up, and swallowing it in a worker thread would + turn that fail-closed answer into a silent hang until the bound expired. + + The call runs inside a copy of the CALLER's context. A new thread starts with an empty one, so + without this the request-scoped ``_timeout_override`` / ``_max_rows_override`` would read as + unset inside the worker and every in-process call would silently fall back to the deployment + defaults. + """ + timeout_s = _resolve_timeout_s() + outcome: dict[str, Any] = {} + ctx = copy_context() + + def call() -> None: + try: + outcome["result"] = ctx.run(executor.execute, sql, creds, profile=profile) + except BaseException as exc: + outcome["error"] = exc + + worker = threading.Thread(target=call, name="agami-bounded-execute", daemon=True) + worker.start() + worker.join(timeout_s + _OUTER_BOUND_SKEW_S) + if worker.is_alive(): + raise _OuterBoundExpired(_OUTLIVED_OUTER_BOUND) + if "error" in outcome: + raise outcome["error"] + return outcome["result"] + + def execute_guarded( sql: str, profile: str, @@ -1572,7 +1662,7 @@ def execute_guarded( * the read-only gate refuses -> ``refused`` carrying that gate's ``Refusal`` * ``_model_safety`` returns a Refusal -> ``refused`` carrying it verbatim * ``_model_safety`` returns an int -> ``refused`` carrying the interim ``model_safety`` - * the per-statement deadline fired -> ``refused`` carrying ``resource_limit`` + * either time bound fired -> ``refused`` carrying ``resource_limit`` * ``executor.execute`` raises -> ``failed`` carrying a classified ``Failure`` * anything else raises -> ``failed``/``other``, generic message, raw to the log * the statement ran -> ``ok`` carrying the ``ExecResult`` @@ -1621,30 +1711,50 @@ def execute_guarded( "query.", )) creds = _load_credentials(profile, org_id or "local") - result = executor.execute(sql, creds, profile=profile) + # Bounded at the CHOKEPOINT, so the limit reaches every executor rather than only the + # built-in one whose engines carry the inner watchdog. See `_execute_bounded` for the + # mechanism and for the leaked worker it costs on expiry. + result = _execute_bounded(executor, sql, creds, profile=profile) # Inside the try on purpose: an executor that returns `None` (or anything else the contract # does not accept) fails the present-iff check in `Envelope.__post_init__`, and that is a # broken adapter, not a reason for the chokepoint to raise at its caller. return _envelope("ok", data=result) - except _ResourceLimit: + except _ResourceLimit as exc: # AHEAD of both handlers below: `_ResourceLimit` is an `ExecutorError` sibling, not a # subclass, but the catch-all would swallow it and report a bound WE imposed as an # unclassified server break. This is the per-statement bound the contract reserves # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a # fix we can honestly name, unlike the supervisor's kill of a child that never returned. # + # ONE handler for both bounds, and so exactly one refusal per call however many layers were + # armed. When the inner watchdog fired, its marker is what arrives here — the outer layer's + # join returned long before its own budget and has nothing to add, so the inner refusal is + # the answer and cannot be overwritten by a later one. + # # The budget is re-resolved rather than carried on the marker: the marker stays a plain # exception, and nothing between the engine call and here can change the env var or the # request-scoped ContextVar the resolver reads, so it is the same number the watchdog used. timeout_s = _resolve_timeout_s() + # The configured number belongs in the detail — it is a deployment setting, not a data + # value, and a bound the caller cannot see is one it cannot plan around. + if isinstance(exc, _OuterBoundExpired): + # The outer layer holds no connection, so it stopped WAITING rather than stopping the + # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the + # driver call and the statement may still be running. The caller is told what we + # actually know, and the number quoted is the bound that actually elapsed. + detail = ( + f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " + "and the query was abandoned." + ) + else: + detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." return _envelope("refused", refusal=refuse( RULE_RESOURCE_LIMIT, - # The configured number belongs here — it is a deployment setting, not a data value, and - # a bound the caller cannot see is one it cannot plan around. The remediation names only - # what would make THIS statement executable: on the served path the caller is an - # assistant with no shell and no deployment, so naming the environment variable would be - # advice it cannot take, addressed to someone who is not reading. - detail=f"The statement ran longer than the {timeout_s}s limit and was cancelled.", + detail=detail, + # The remediation names only what would make THIS statement executable: on the served + # path the caller is an assistant with no shell and no deployment, so naming the + # environment variable would be advice it cannot take, addressed to someone who is not + # reading. remediation="Narrow the time range, reduce the grouping, or add a selective filter, " "then run it again.", )) diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index 1bbbd76..82b3a99 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -1414,13 +1414,24 @@ def tool_execute_sql(args: dict[str, Any]) -> str: # materializes the whole result — not a client-side trim after the fact. cmd += ["--max-rows", str(max_rows)] + # The supervisor's bound is the OUTERMOST of the four time bounds, and it is DERIVED from the + # same resolver every inner layer reads rather than being a number of its own. A fixed 240s + # inverted that order for any statement budget approaching it: the supervisor fired FIRST, so a + # statement we could have cancelled and refused precisely came back instead as a + # `failed`/`timeout` naming nothing the caller can act on. Imported lazily for the same + # reason `_run_in_process` does it. The child inherits `os.environ` (no `env=` below), so it + # resolves the identical `AGAMI_SQL_TIMEOUT_S` this call just read. + import execute_sql + + supervisor_timeout_s = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S + started = time.monotonic() try: proc = subprocess.run( cmd, capture_output=True, text=True, - timeout=240, + timeout=supervisor_timeout_s, ) except subprocess.TimeoutExpired: # A `failed`/`timeout`, NOT a `refused`/`resource_limit` — and the reason is what we can diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 1f396ff..d53ef2d 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -63,7 +63,7 @@ import urllib.parse import uuid from collections.abc import Callable, Iterator -from contextvars import ContextVar +from contextvars import ContextVar, copy_context from dataclasses import asdict, dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -1082,6 +1082,26 @@ def _resolve_row_cap() -> int: # backstop for the one case the watchdog cannot cover — our process dying with a statement in # flight, which would otherwise leave the engine scanning for nobody. _NATIVE_BOUND_SKEW_S = 5 +# How far behind the watchdog the OUTER bound sits — the one `execute_guarded` puts around +# `executor.execute` itself, so an INJECTED executor is bounded too. The four bounds are one ordered +# family resolved from the same budget, innermost first: +# +# watchdog (timeout_s) < native (+5) < outer (+10) < supervisor (+60) +# +# and the order is the contract. The inner watchdog must always win, because it is the only layer +# that can attribute the stop to the statement and hand the caller the precise refusal; every layer +# behind it is a backstop for a failure the layer in front of it cannot see. Collapse the order and +# a bound we imposed comes back as something else — a native server-side kill reads as an ordinary +# database error, and a supervisor kill reads as `failed`/`timeout` naming nothing the caller can +# act on. Both are strictly worse answers to the same event. +_OUTER_BOUND_SKEW_S = 10 +# The OUTERMOST bound: how far behind the watchdog the subprocess supervisor in `tools` stops waiting +# for a forked child. It lives here, next to its three siblings, so the whole family is derived from +# one resolver and one place — the supervisor was a hardcoded 240s, which for any statement budget +# approaching it became the FIRST bound to fire and inverted the order. Its slack is large because it +# bounds a whole process (interpreter start, model load, credential resolution, connect) rather than +# a statement. +_SUPERVISOR_SKEW_S = 60 # Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the # HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the # context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one @@ -1124,10 +1144,24 @@ class _ResourceLimit(Exception): boundary.""" +class _OuterBoundExpired(_ResourceLimit): + """Raised when the OUTER bound expired: the executor never returned and we stopped waiting. + + A subclass, so the one `except _ResourceLimit` handler in `execute_guarded` produces exactly one + refusal for either layer and no second one can be minted. It stays distinguishable because the + two events are not the same event: the watchdog CANCELLED a statement it holds a connection to, + while this layer only abandoned its wait — so the sentence the caller reads differs, and saying + "cancelled" here would be false. + """ + + # The marker's message, single-sourced because every engine raises it. It is diagnostic text, not # caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own # detail, so nothing a caller reads comes from here. _OUTLIVED_BUDGET = "the statement outlived its per-statement budget" +# The outer marker's message, and diagnostic in the same way — it names the executor rather than the +# statement, because at this layer the statement is not the thing we observed. +_OUTLIVED_OUTER_BOUND = "the executor outlived the outer bound around it" @contextlib.contextmanager @@ -1548,6 +1582,62 @@ def _envelope( ) +def _execute_bounded( + executor: Executor, sql: str, creds: dict[str, str], *, profile: str +) -> ExecResult: + """Run ``executor.execute`` under the OUTER bound and return what it returned. + + This is the layer that makes "no executor escapes the limit" true. The inner deadline lives + inside the BUILT-IN executor's engine functions, so before this an INJECTED executor — the + hosted connection-reuse path — ran with no per-statement bound at all, and even on the built-in + path BigQuery has no watchdog (it has no connection object to cancel), so a client-side stall + there was unbounded too. Both are bounded here, and deliberately by the same layer: applying + this only to non-built-in executors would leave the BigQuery hole open while looking closed. + + **The mechanism is a worker thread, because this layer holds nothing it can cancel.** An + arbitrary ``Executor`` exposes one method and no connection, so the only thing this code owns is + its own WAIT — it starts the call on a daemon thread and joins with the budget. + + **The cost is a leaked worker, and it is real.** On expiry the thread is still inside the + driver call, and it stays there: nothing here cancels it, and it may hold a database connection + (and, on the server side, a running statement) until that call returns on its own. It is a + daemon so it cannot hold the interpreter open at exit, and that is the whole of the mitigation. + This is a bound on how long a CALLER waits, not a promise that the work stopped — which is + exactly why the inner watchdog, the layer that really can cancel, is set to fire first. + + Exceptions cross the thread boundary with their ORIGINAL type, re-raised here. That is what + keeps the handlers in ``execute_guarded`` correct: ``_ResourceLimit`` still reaches the refusal + branch and ``ExecutorError`` still reaches the classified one, rather than every executor + failure arriving as a worker that finished having produced nothing. ``BaseException`` is caught + rather than ``Exception`` for the same reason: a driver's deep ``sys.exit`` raises ``SystemExit``, + which ``tools._run_in_process`` nets one layer up, and swallowing it in a worker thread would + turn that fail-closed answer into a silent hang until the bound expired. + + The call runs inside a copy of the CALLER's context. A new thread starts with an empty one, so + without this the request-scoped ``_timeout_override`` / ``_max_rows_override`` would read as + unset inside the worker and every in-process call would silently fall back to the deployment + defaults. + """ + timeout_s = _resolve_timeout_s() + outcome: dict[str, Any] = {} + ctx = copy_context() + + def call() -> None: + try: + outcome["result"] = ctx.run(executor.execute, sql, creds, profile=profile) + except BaseException as exc: + outcome["error"] = exc + + worker = threading.Thread(target=call, name="agami-bounded-execute", daemon=True) + worker.start() + worker.join(timeout_s + _OUTER_BOUND_SKEW_S) + if worker.is_alive(): + raise _OuterBoundExpired(_OUTLIVED_OUTER_BOUND) + if "error" in outcome: + raise outcome["error"] + return outcome["result"] + + def execute_guarded( sql: str, profile: str, @@ -1572,7 +1662,7 @@ def execute_guarded( * the read-only gate refuses -> ``refused`` carrying that gate's ``Refusal`` * ``_model_safety`` returns a Refusal -> ``refused`` carrying it verbatim * ``_model_safety`` returns an int -> ``refused`` carrying the interim ``model_safety`` - * the per-statement deadline fired -> ``refused`` carrying ``resource_limit`` + * either time bound fired -> ``refused`` carrying ``resource_limit`` * ``executor.execute`` raises -> ``failed`` carrying a classified ``Failure`` * anything else raises -> ``failed``/``other``, generic message, raw to the log * the statement ran -> ``ok`` carrying the ``ExecResult`` @@ -1621,30 +1711,50 @@ def execute_guarded( "query.", )) creds = _load_credentials(profile, org_id or "local") - result = executor.execute(sql, creds, profile=profile) + # Bounded at the CHOKEPOINT, so the limit reaches every executor rather than only the + # built-in one whose engines carry the inner watchdog. See `_execute_bounded` for the + # mechanism and for the leaked worker it costs on expiry. + result = _execute_bounded(executor, sql, creds, profile=profile) # Inside the try on purpose: an executor that returns `None` (or anything else the contract # does not accept) fails the present-iff check in `Envelope.__post_init__`, and that is a # broken adapter, not a reason for the chokepoint to raise at its caller. return _envelope("ok", data=result) - except _ResourceLimit: + except _ResourceLimit as exc: # AHEAD of both handlers below: `_ResourceLimit` is an `ExecutorError` sibling, not a # subclass, but the catch-all would swallow it and report a bound WE imposed as an # unclassified server break. This is the per-statement bound the contract reserves # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a # fix we can honestly name, unlike the supervisor's kill of a child that never returned. # + # ONE handler for both bounds, and so exactly one refusal per call however many layers were + # armed. When the inner watchdog fired, its marker is what arrives here — the outer layer's + # join returned long before its own budget and has nothing to add, so the inner refusal is + # the answer and cannot be overwritten by a later one. + # # The budget is re-resolved rather than carried on the marker: the marker stays a plain # exception, and nothing between the engine call and here can change the env var or the # request-scoped ContextVar the resolver reads, so it is the same number the watchdog used. timeout_s = _resolve_timeout_s() + # The configured number belongs in the detail — it is a deployment setting, not a data + # value, and a bound the caller cannot see is one it cannot plan around. + if isinstance(exc, _OuterBoundExpired): + # The outer layer holds no connection, so it stopped WAITING rather than stopping the + # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the + # driver call and the statement may still be running. The caller is told what we + # actually know, and the number quoted is the bound that actually elapsed. + detail = ( + f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " + "and the query was abandoned." + ) + else: + detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." return _envelope("refused", refusal=refuse( RULE_RESOURCE_LIMIT, - # The configured number belongs here — it is a deployment setting, not a data value, and - # a bound the caller cannot see is one it cannot plan around. The remediation names only - # what would make THIS statement executable: on the served path the caller is an - # assistant with no shell and no deployment, so naming the environment variable would be - # advice it cannot take, addressed to someone who is not reading. - detail=f"The statement ran longer than the {timeout_s}s limit and was cancelled.", + detail=detail, + # The remediation names only what would make THIS statement executable: on the served + # path the caller is an assistant with no shell and no deployment, so naming the + # environment variable would be advice it cannot take, addressed to someone who is not + # reading. remediation="Narrow the time range, reduce the grouping, or add a selective filter, " "then run it again.", )) diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index 53f41b7..50cae3d 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -19,6 +19,13 @@ mean an unlucky query gets told to narrow itself when nothing timed out. `_deadline` sets its Event *before* the cancel lands, so "did WE stop this?" is answerable without inference — and it is asserted here in both directions. + +The last section is the OUTER bound, which is what makes the limit apply to every executor rather +than to the built-in one alone. An injected executor (the hosted connection-reuse path) carries none +of the engine watchdogs, and even the built-in BigQuery path has no cancel to arm — so +`execute_guarded` bounds `executor.execute` itself, on a worker thread it can stop waiting for. The +four bounds are one ordered family (watchdog < native < outer < supervisor) resolved from one budget, +and the order is asserted as a single fact so no future change can let them drift apart. """ from __future__ import annotations @@ -26,7 +33,9 @@ import contextlib import json import logging +import os import sqlite3 +import subprocess import sys import threading import time @@ -51,10 +60,14 @@ @pytest.fixture(autouse=True) def _reset_override(): - # _timeout_override is a request-scoped ContextVar; isolate every test from it. + # Both overrides are request-scoped ContextVars; isolate every test from them. The row cap is + # reset here too because the outer bound copies the caller's whole context into its worker, so a + # cap left set by one test is now visible to the next test's executor as well. execute_sql._timeout_override.set(None) + execute_sql._max_rows_override.set(None) yield execute_sql._timeout_override.set(None) + execute_sql._max_rows_override.set(None) @pytest.fixture(autouse=True) @@ -63,6 +76,21 @@ def _clear_env(monkeypatch): monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) +@pytest.fixture(autouse=True) +def _no_injected_executor(): + """`_INJECTED_EXECUTOR` is a process global, and the tests below install one both directly and + (via `create_app`) as a side effect of building the HTTP app. Reset it around every test so + neither leaks into the next one and quietly moves it onto the other execution path.""" + try: + import tools + except Exception: + yield + return + tools.set_injected_executor(None) + yield + tools.set_injected_executor(None) + + # -------------------------------------------------------------------------------------------- # _resolve_timeout_s # -------------------------------------------------------------------------------------------- @@ -563,6 +591,10 @@ def test_an_error_that_merely_arrives_late_is_still_a_failure(warehouse, fake_sq # The refusal is recorded # -------------------------------------------------------------------------------------------- +# A neutral hostname and a throwaway secret, only so the HTTP surface below can mint a token. +_BASE_URL = "https://your-host.example.com" +_SIGNING_SECRET = "x" * 40 + @pytest.fixture def audited(tmp_path, monkeypatch): @@ -607,6 +639,10 @@ def audited(tmp_path, monkeypatch): monkeypatch.delenv("AGAMI_ORG_ID", raising=False) monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "artifacts")) monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", f"sqlite:///{path}") + # The HTTP surface needs both to mint and verify its bearer token; the stdio surface ignores + # them. Set here so one install serves both transports. + monkeypatch.setenv("PUBLIC_BASE_URL", _BASE_URL) + monkeypatch.setenv("AGAMI_SIGNING_SECRET", _SIGNING_SECRET) return SimpleNamespace(app_db=app_db) @@ -645,6 +681,374 @@ def test_the_refusal_is_written_to_the_audit_trail(audited): assert row["reason"] == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] +# -------------------------------------------------------------------------------------------- +# The outer bound — the limit reaches every executor, not only the built-in one +# -------------------------------------------------------------------------------------------- +# +# The engine watchdogs live INSIDE the built-in executor, so an injected one — which is what the +# hosted connection-reuse path supplies — ran with no per-statement bound at all, and BigQuery has no +# watchdog even on the built-in path. `execute_guarded` therefore bounds `executor.execute` itself. +# That layer holds no connection and can cancel nothing; it runs the call on a daemon worker and +# stops WAITING, which is why it must be the outer one and the watchdog must always win. + + +class _BlockingExecutor: + """A `ports.Executor` whose `execute` never returns on its own — the shape the outer bound + exists for. Nothing about it is cancellable from outside, which is exactly the point: an + arbitrary injected executor exposes one method and no handle on the work behind it.""" + + def __init__(self) -> None: + self.calls = 0 + self.entered = threading.Event() + self.release = threading.Event() + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: + self.calls += 1 + self.entered.set() + # Released by the test, so the worker this deliberately leaks does not outlive the test that + # made it. A real one would sit here until its driver call returned on its own. + self.release.wait(10) + return execute_sql.ExecResult(columns=["c"], rows=[(1,)], truncated=False) + + +def test_an_injected_executor_that_never_returns_is_still_bounded(warehouse, monkeypatch): + """The headline for this layer: an executor with no watchdog of its own, no cancel and no + intention of returning still yields a `resource_limit` refusal at a bound we set. + + Without the outer layer this call blocks in the executor and comes back `ok` with its rows, so + the assertion is on the WHOLE outcome — that it returned at all, in time, and as a refusal. + + The skew is patched to zero so the bound under test is the 1s budget rather than 11s of wall + clock. The +10 value it normally carries is not skipped: it is pinned by the ordering test below, + which needs no clock at all. + """ + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + execute_sql._timeout_override.set(_BUDGET_S) + executor = _BlockingExecutor() + + started = time.monotonic() + try: + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=executor, no_safety=True, + ) + elapsed = time.monotonic() - started + finally: + executor.release.set() + + assert executor.entered.is_set() # the executor really ran; this is not a gate refusing early + assert elapsed < 10, f"the call waited {elapsed:.1f}s on a {_BUDGET_S}s bound" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + # A bound WE imposed is a refusal, not a failure — and it carries no data, like every refusal. + assert env.data is None and env.failure is None + + +def test_the_outer_refusal_says_what_actually_happened(warehouse, monkeypatch): + """It must not claim the statement was cancelled. Nothing was: the worker is still inside the + driver call, may still hold a connection, and the statement may still be running on the server. + The sentence the caller reads names what we observed — the executor did not come back — and + quotes the bound that actually elapsed rather than the inner budget.""" + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + execute_sql._timeout_override.set(_BUDGET_S) + executor = _BlockingExecutor() + + try: + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=executor, no_safety=True, + ) + finally: + executor.release.set() + + assert "cancelled" not in env.refusal.detail, env.refusal.detail + assert "did not return" in env.refusal.detail, env.refusal.detail + # The remediation is the same one the inner bound gives: the fix that makes THIS statement + # runnable, addressed to a caller with no shell and no deployment. + authored = f"{env.refusal.detail} {env.refusal.remediation}" + assert "AGAMI_" not in authored, authored + + +def test_the_inner_refusal_is_the_one_the_caller_gets(warehouse): + """When both layers are armed the INNER one wins, and the outer neither overwrites its refusal + nor mints a second. + + That ordering is what buys the caller a precise answer: only the watchdog holds the connection + it cancelled, so only it can say the STATEMENT was stopped. A real cancel is driven here (no + stub, no patched skew) with the outer bound left at its real +10, so the outer layer is genuinely + armed and genuinely loses the race. The detail is the assertion, because it is the one part of + the envelope the two layers word differently. + """ + execute_sql._timeout_override.set(_BUDGET_S) + + started = time.monotonic() + env = _guarded(_RUNAWAY_SQL) + elapsed = time.monotonic() - started + + assert elapsed < 20, f"the statement ran {elapsed:.1f}s against a {_BUDGET_S}s budget" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert f"{_BUDGET_S}s" in env.refusal.detail, env.refusal.detail # the inner budget + assert "cancelled" in env.refusal.detail, env.refusal.detail # the inner sentence + assert "did not return" not in env.refusal.detail, env.refusal.detail # not the outer one + + +def test_the_four_bounds_are_ordered_from_one_resolved_budget(monkeypatch): + """One fact, asserted once: the four time bounds are derived from a single resolved budget and + stand in a fixed order — watchdog < native (+5) < outer (+10) < supervisor (+60). + + Asserted together rather than one skew per test because the ORDER is the property, and four + separate assertions would each stay green while the family drifted apart. Each layer behind the + watchdog is a backstop for something the layer in front cannot see, and every inversion trades a + precise refusal for a vaguer answer to the same event: a native kill reads as a database error, + an outer expiry cannot say the statement stopped, and a supervisor kill cannot say what hung. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") # distinctive: a hard-coded 30 cannot pass + timeout_s = execute_sql._resolve_timeout_s() + + watchdog = timeout_s + native = timeout_s + execute_sql._NATIVE_BOUND_SKEW_S + outer = timeout_s + execute_sql._OUTER_BOUND_SKEW_S + supervisor = timeout_s + execute_sql._SUPERVISOR_SKEW_S + + assert watchdog < native < outer < supervisor + assert (watchdog, native, outer, supervisor) == (45, 50, 55, 105) + + +class _RaisingExecutor: + """A `ports.Executor` that raises a given exception from inside the worker thread.""" + + def __init__(self, exc: BaseException) -> None: + self._exc = exc + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: + raise self._exc + + +def test_an_executor_error_keeps_its_type_across_the_worker(warehouse): + """The subtle half of running the call off-thread: an exception raised in the worker has to reach + the caller with its ORIGINAL type, or every handler in `execute_guarded` stops matching. + + `ExecutorError` is the one that would fail loudest — its handler is what turns a driver error + into a classified `failed` envelope carrying the message this module authored. Caught as anything + else it would land in the catch-all and the caller would get the generic string instead, so the + relayed message is asserted rather than just the status. + """ + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_RaisingExecutor(execute_sql.ExecutorError("no such column: nope", code=5)), + no_safety=True, + ) + + assert env.status == "failed" + assert env.failure.kind == "syntax" # the classified branch, i.e. `except ExecutorError` matched + assert env.failure.message == "no such column: nope" + assert env.failure.message != execute_sql.UNEXPECTED_FAILURE_MESSAGE + + +def test_a_plain_exception_still_reaches_the_catch_all_across_the_worker(warehouse): + """And the other direction: an exception nobody classified is still an unanticipated break, told + to the caller as the one fixed string. A worker that swallowed it would leave the call looking + like it produced nothing at all.""" + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_RaisingExecutor(RuntimeError("a driver detail nobody has vetted")), + no_safety=True, + ) + + assert env.status == "failed" + assert env.failure.kind == "other" + assert env.failure.message == execute_sql.UNEXPECTED_FAILURE_MESSAGE + assert "driver detail" not in env.failure.message # the raw text goes to the log, not the caller + + +class _ContextProbeExecutor: + """Records the thread it ran on and what the two request-scoped ContextVars read there.""" + + def __init__(self) -> None: + self.thread: threading.Thread | None = None + self.seen_timeout: object = "not read" + self.seen_rows: object = "not read" + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: + self.thread = threading.current_thread() + self.seen_timeout = execute_sql._timeout_override.get() + self.seen_rows = execute_sql._max_rows_override.get() + return execute_sql.ExecResult(columns=["c"], rows=[(1,)], truncated=False) + + +def test_the_request_scoped_overrides_are_visible_inside_the_worker(warehouse): + """A new thread starts with an EMPTY context, so without an explicit copy the caller's per-call + budget and row cap would read as unset inside the worker and every in-process call would quietly + fall back to the deployment defaults — a bug with no symptom other than the wrong numbers. + + The thread identity is asserted first, because it is what makes the rest of this test mean + anything: if the call ever stopped running off-thread, the ContextVar assertions would pass + trivially and the copy could be deleted with the suite still green. + """ + execute_sql._timeout_override.set(7) + execute_sql._max_rows_override.set(5) + probe = _ContextProbeExecutor() + + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=probe, no_safety=True, + ) + + assert env.status == "ok" + assert probe.thread is not threading.current_thread() # it really ran on a worker + assert probe.seen_timeout == 7 + assert probe.seen_rows == 5 + + +# -------------------------------------------------------------------------------------------- +# The supervisor's bound, derived rather than fixed +# -------------------------------------------------------------------------------------------- + + +def test_the_supervisor_bound_is_derived_from_the_statement_budget(warehouse, monkeypatch): + """The fork path's supervisor was a hardcoded 240s, which quietly became the FIRST bound to fire + on any deployment configuring a larger statement budget — turning a refusal that could name the + statement into a `failed`/`timeout` that names nothing. It is now `budget + 60`, from the same + resolver every inner layer reads. + + Asserted at a RAISED budget, which is the case the fixed number got wrong: at 300s the supervisor + must be 360, and 240 would have killed the child a full minute before its own statement bound. + The computed value is read off the `subprocess.run` call rather than waited out, for the obvious + reason. + """ + import tools + + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "300") + monkeypatch.setattr(tools, "resolve_profile", lambda ds: PROFILE) + + class _Proc: + returncode = 0 + stdout = "id\r\n1\r\n" + stderr = "" + + captured: dict = {} + + def _fake_run(cmd, **kwargs): + captured.update(kwargs) + return _Proc() + + monkeypatch.setattr(tools.subprocess, "run", _fake_run) + tools.set_injected_executor(None) # the fork path, which is the one with a supervisor + + tools.tool_execute_sql({"sql": "SELECT id FROM orders", "datasource": PROFILE}) + + assert captured["timeout"] == 300 + execute_sql._SUPERVISOR_SKEW_S == 360 + assert captured["timeout"] > 300, "the supervisor must fire AFTER the statement bound, not before" + assert captured["timeout"] != 240, "the fixed bound this replaces" + # And nothing was passed that would stop the child inheriting the same configuration. + assert "env" not in captured + + +def test_the_supervisors_verdict_is_unchanged(warehouse, monkeypatch): + """Only the DURATION moved. A child that never came back is still a `failed`/`timeout` and never + a `resource_limit`: the bound is ours, but all it tells us is that the child stopped responding — + it may have hung in connect, in credential resolution or in loading the model, and on any of + those "narrow the query" points at the wrong thing.""" + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: PROFILE) + + def _timed_out(cmd, **kwargs): + raise subprocess.TimeoutExpired(cmd, kwargs.get("timeout")) + + monkeypatch.setattr(tools.subprocess, "run", _timed_out) + tools.set_injected_executor(None) + + body = json.loads(tools.tool_execute_sql({"sql": "SELECT id FROM orders", + "datasource": PROFILE})) + + assert body["status"] == "failed" + assert body["failure"]["kind"] == "timeout" + + +# -------------------------------------------------------------------------------------------- +# Both surfaces refuse identically +# -------------------------------------------------------------------------------------------- + + +def _stdio_refusal(sql: str) -> dict: + """`python -m mcp_harness` over JSON-RPC on stdin — the transport a desktop client launches, and + the one that forks `python -m execute_sql`. The child inherits this process's environment, so it + resolves the same budget, model and warehouse; nothing on this path is stubbed.""" + messages = [ + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + {"jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "execute_sql", "arguments": {"sql": sql, "datasource": PROFILE, + "raw_query": "how many"}}}, + ] + proc = subprocess.run( + [sys.executable, "-m", "mcp_harness"], + input="".join(json.dumps(m) + "\n" for m in messages), + capture_output=True, text=True, timeout=180, env={**os.environ}, + ) + replies = { + m.get("id"): m + for m in (json.loads(line) for line in proc.stdout.splitlines() if line.strip()) + } + assert 2 in replies, proc.stderr + return json.loads(replies[2]["result"]["content"][0]["text"]) + + +def _http_refusal(sql: str) -> dict: + """The authenticated HTTP transport, which runs execution IN-PROCESS through `create_app()`'s + default adapters — the other of the two execution paths.""" + import mcp_http + import tools + from oauth_server import issue_jwt + from starlette.testclient import TestClient + + headers = { + "Authorization": f"Bearer {issue_jwt('jordan@example.com')}", + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + with TestClient(mcp_http.create_app()) as client: + assert tools._INJECTED_EXECUTOR is not None, ( + "create_app() no longer injects an executor, so this surface is now the fork path and " + "the in-process path this test believes it covers is uncovered" + ) + init = client.post("/mcp", headers=headers, json={ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}, + "clientInfo": {"name": "t", "version": "1"}}}) + session = init.headers.get("mcp-session-id") + headers2 = {**headers, **({"mcp-session-id": session} if session else {})} + client.post("/mcp", headers=headers2, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}) + resp = client.post("/mcp", headers=headers2, json={ + "jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "execute_sql", "arguments": {"sql": sql, "datasource": PROFILE, + "raw_query": "how many"}}}) + assert resp.status_code == 200, resp.text + return json.loads(resp.json()["result"]["content"][0]["text"]) + + +def test_both_surfaces_refuse_a_runaway_statement_identically(audited, monkeypatch): + """The criterion in one test: the limit applies the same on stdio and on HTTP, and no executor + escapes it. + + The two transports run DIFFERENT execution paths — stdio forks a child and reads its stderr back, + HTTP runs the built-in executor in-process behind the outer bound — so "identically" is asserted + on the whole refusal, not merely on the status. A difference in any field would mean the caller's + answer depends on which door it came through. + """ + pytest.importorskip("starlette") + pytest.importorskip("mcp") + pytest.importorskip("jwt") + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + stdio = _stdio_refusal(_RUNAWAY_SQL) + http = _http_refusal(_RUNAWAY_SQL) + + assert stdio["status"] == http["status"] == "refused", (stdio, http) + assert stdio["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT + assert stdio["refusal"] == http["refusal"] + + # -------------------------------------------------------------------------------------------- # Every engine, under one table # -------------------------------------------------------------------------------------------- From 9da6925f8aa4baaa2c4f9cb1f01c2d295de5b113 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 16:26:09 -0700 Subject: [PATCH 05/10] test(executor): prove the cancel reaches the server, and document the budget (S5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-process tests prove the refusal; none of them can prove the cancel. A client that merely stopped waiting produces the identical Envelope while the backend keeps scanning for nobody, so the security claim has to be asserted against the SERVER's view. `test_postgres_timeout_integration.py` does that on a live Postgres, from a second connection: the statement is observed running in `pg_stat_activity`, and after the refusal it is gone. It also proves the native backstop landed by disarming our watchdog and letting Postgres kill the statement itself — "canceling statement due to statement timeout" is wording a client cancel cannot produce. The runaway is three narrow cross-joined series rather than one wide one: a function scan materializes into a tuplestore, so the obvious `generate_series(1, 1e11)` burns time by filling the server's temp tablespace, and a test that leaves the disk full has done more damage than the bug it checked for. Opt-in on the existing gate (skips unless AGAMI_IT_PG_PASSWORD is set), so CI without a database is unaffected. `AGAMI_SQL_TIMEOUT_S` reaches both env examples — it was the one bound an operator had no way to discover. AGAMI_SQL_MAX_ROWS is left undocumented on purpose: its truncation behaviour is being changed by separate work. The four ACE-038 comment markers in execute_sql.py (and three more in tools.py) sat on row-cap and result-transfer code that belongs to a different spec. The prose stays; the ids go. Spec ids belong in the PR trailer. --- deploy/agami.env.example | 7 + packages/agami-core/src/execute_sql.py | 8 +- packages/agami-core/src/tools.py | 6 +- plugins/agami/lib/execute_sql.py | 8 +- .../agami-deploy/bundle/agami.env.example | 7 + tests/test_ace038_timeout.py | 5 + tests/test_postgres_timeout_integration.py | 268 ++++++++++++++++++ 7 files changed, 298 insertions(+), 11 deletions(-) create mode 100644 tests/test_postgres_timeout_integration.py diff --git a/deploy/agami.env.example b/deploy/agami.env.example index 2d23dd0..bce775d 100644 --- a/deploy/agami.env.example +++ b/deploy/agami.env.example @@ -42,6 +42,13 @@ AGAMI_ADMIN_PASSWORD= # Override the bundled-Postgres password if you like (it's never exposed, so the default is fine). # POSTGRES_PASSWORD= +# Optional: how long ONE statement may run, in whole SECONDS (unset = 30). This is the deployment's +# availability budget — a query that outruns it is cancelled on the warehouse and the caller is told +# to narrow it. Raise it if your warehouse is slow and your users would rather wait; lower it to keep +# a runaway query from holding a connection. The executor's other time bounds are derived from this +# one, so this is the single number to tune. +# AGAMI_SQL_TIMEOUT_S=30 # per-statement budget (default 30 seconds) + # Optional free registration key (best-effort; absent = full run). # AGAMI_LICENSE_KEY= diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index d53ef2d..bd267e8 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -508,7 +508,7 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: # A server-side (named) cursor so the row cap bounds TRANSFER, not just what we write: # psycopg2's default client-side cursor buffers the ENTIRE result before we can fetchmany, # so a runaway result would still be pulled whole. The named cursor streams from the - # server in bounded batches (ACE-038). Read-only SELECTs (the only thing the guard admits) + # server in bounded batches. Read-only SELECTs (the only thing the guard admits) # are exactly what a server-side cursor supports. cur = conn.cursor(name="agami_bounded") cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window @@ -1049,7 +1049,7 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: return result -_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation (ACE-038) +_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation # Per-call cap from --max-rows (ACE-044). A ContextVar, not a plain global, so it is REQUEST-SCOPED # once the HTTP server runs execution in-process (ACE-028): concurrent handlers run in worker threads # (`run_blocking` copies the context per call, like `_current_org_ctx`), so each request's cap is @@ -1195,7 +1195,7 @@ def fire() -> None: def _flag_truncated(cap: int) -> None: """Signal a bounded-fetch truncation to the caller — a non-error `{"truncated": …}` marker on stderr (distinct from the guards' `{"error": …}`), so a truncated result is never mistaken for a - complete one (ACE-038/044). Shared by every engine's materialization path. One write so the + complete one (ACE-044). Shared by every engine's materialization path. One write so the marker is always a single line, even if other notices surround it.""" sys.stderr.write(json.dumps({"truncated": {"row_cap": cap}}) + "\n") @@ -1208,7 +1208,7 @@ def _collect_cursor(cur: Any) -> ExecResult: executor path share, so the row cap is enforced once, identically, for every caller. Fetch FIRST, then read ``cur.description``: a psycopg2 **server-side (named) cursor** — which the - Postgres/Redshift path uses to bound transfer (ACE-038) — reports ``description is None`` until the + Postgres/Redshift path uses to bound transfer — reports ``description is None`` until the first fetch, so reading it beforehand would drop EVERY row of a real Postgres result. Client-side cursors (sqlite/mysql/…) set ``description`` at execute, so fetch-first is equally correct there.""" cap = _resolve_row_cap() diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index 82b3a99..4b92f1e 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -1063,7 +1063,7 @@ def _child_failure_message(returncode: int, stderr: str | None) -> str: def _executor_truncated(stderr: str | None) -> bool: - """True if execute_sql flagged a bounded-fetch truncation (ACE-038/044). The executor emits a + """True if execute_sql flagged a bounded-fetch truncation (ACE-044). The executor emits a non-error `{"truncated": {"row_cap": N}}` line on stderr alongside any other notices; scan for it.""" for line in (stderr or "").splitlines(): line = line.strip() @@ -1205,7 +1205,7 @@ def _emit( columns = list(env.data.columns) rows = [["" if v is None else str(v) for v in row] for row in env.data.rows] truncated = env.data.truncated - # Backstop only: the executor already caps at the source (ACE-038/044) and flags it. This + # Backstop only: the executor already caps at the source (ACE-044) and flags it. This # catches a result that slipped past that bound, and marks it truncated rather than # presenting a trimmed result as complete. if max_rows is not None and len(rows) > max_rows: @@ -1479,7 +1479,7 @@ def tool_execute_sql(args: dict[str, Any]) -> str: env, sql=sql, execution_ms=execution_ms, profile=profile, args=args, ) - # Parse the RFC-4180 CSV emitted on stdout. The executor caps at the source (ACE-038/044) and + # Parse the RFC-4180 CSV emitted on stdout. The executor caps at the source (ACE-044) and # flags it on stderr; carry that flag so a truncated result is never presented as complete. The # `max_rows` backstop is applied by `_emit`, the same one the in-process path gets. reader = csv.reader(io.StringIO(proc.stdout)) diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index d53ef2d..bd267e8 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -508,7 +508,7 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: # A server-side (named) cursor so the row cap bounds TRANSFER, not just what we write: # psycopg2's default client-side cursor buffers the ENTIRE result before we can fetchmany, # so a runaway result would still be pulled whole. The named cursor streams from the - # server in bounded batches (ACE-038). Read-only SELECTs (the only thing the guard admits) + # server in bounded batches. Read-only SELECTs (the only thing the guard admits) # are exactly what a server-side cursor supports. cur = conn.cursor(name="agami_bounded") cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window @@ -1049,7 +1049,7 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: return result -_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation (ACE-038) +_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation # Per-call cap from --max-rows (ACE-044). A ContextVar, not a plain global, so it is REQUEST-SCOPED # once the HTTP server runs execution in-process (ACE-028): concurrent handlers run in worker threads # (`run_blocking` copies the context per call, like `_current_org_ctx`), so each request's cap is @@ -1195,7 +1195,7 @@ def fire() -> None: def _flag_truncated(cap: int) -> None: """Signal a bounded-fetch truncation to the caller — a non-error `{"truncated": …}` marker on stderr (distinct from the guards' `{"error": …}`), so a truncated result is never mistaken for a - complete one (ACE-038/044). Shared by every engine's materialization path. One write so the + complete one (ACE-044). Shared by every engine's materialization path. One write so the marker is always a single line, even if other notices surround it.""" sys.stderr.write(json.dumps({"truncated": {"row_cap": cap}}) + "\n") @@ -1208,7 +1208,7 @@ def _collect_cursor(cur: Any) -> ExecResult: executor path share, so the row cap is enforced once, identically, for every caller. Fetch FIRST, then read ``cur.description``: a psycopg2 **server-side (named) cursor** — which the - Postgres/Redshift path uses to bound transfer (ACE-038) — reports ``description is None`` until the + Postgres/Redshift path uses to bound transfer — reports ``description is None`` until the first fetch, so reading it beforehand would drop EVERY row of a real Postgres result. Client-side cursors (sqlite/mysql/…) set ``description`` at execute, so fetch-first is equally correct there.""" cap = _resolve_row_cap() diff --git a/plugins/agami/skills/agami-deploy/bundle/agami.env.example b/plugins/agami/skills/agami-deploy/bundle/agami.env.example index 7d7bc23..8846644 100644 --- a/plugins/agami/skills/agami-deploy/bundle/agami.env.example +++ b/plugins/agami/skills/agami-deploy/bundle/agami.env.example @@ -37,6 +37,13 @@ AGAMI_ADMIN_PASSWORD= # Override the bundled-Postgres password if you like (it's never exposed, so the default is fine). # POSTGRES_PASSWORD= +# Optional: how long ONE statement may run, in whole SECONDS (unset = 30). This is the deployment's +# availability budget — a query that outruns it is cancelled on the warehouse and the caller is told +# to narrow it. Raise it if your warehouse is slow and your users would rather wait; lower it to keep +# a runaway query from holding a connection. The executor's other time bounds are derived from this +# one, so this is the single number to tune. +# AGAMI_SQL_TIMEOUT_S=30 # per-statement budget (default 30 seconds) + # Only for the `tunnel` profile: # CLOUDFLARE_TUNNEL_TOKEN= diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index 50cae3d..d6c8dab 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -13,6 +13,11 @@ `resource_limit` — with no partial data, a detail that quotes the configured budget, and a remediation addressed to whoever can actually act on it. +What no test here can show is that the cancel reached a SERVER: SQLite is in-process, and on a +client/server engine an abandoned statement produces the identical Envelope while the backend keeps +running. That assertion needs a live database and a second connection, and lives in +`test_postgres_timeout_integration.py`, which skips unless one is configured. + **The classification is the FLAG, and only the flag.** A cancelled SQLite statement raises `OperationalError("interrupted")`, so neither the error text nor the elapsed clock can be the test: both are properties an ordinary database error can have by coincidence, and reading either one would diff --git a/tests/test_postgres_timeout_integration.py b/tests/test_postgres_timeout_integration.py new file mode 100644 index 0000000..63e8348 --- /dev/null +++ b/tests/test_postgres_timeout_integration.py @@ -0,0 +1,268 @@ +"""Integration guard: the per-statement timeout on a LIVE Postgres, asserted against the SERVER's +view of what happened rather than the client's. + +The in-process tests prove the refusal; they cannot prove the cancel. A client that simply stopped +waiting — abandoned the statement and closed nothing — produces the identical Envelope while the +backend keeps scanning for nobody. So the assertion here is made from a SECOND connection: the +statement is observed running in `pg_stat_activity`, and after the refusal it is GONE from +`pg_stat_activity`. That is the difference between a bound on the caller and a bound on the +database, and it is the whole security value of the feature. + +The second test proves the other half — that `SET LOCAL statement_timeout` really landed on the +session. It disarms our own watchdog and lets the statement run past the native bound; Postgres then +kills it itself, and says so in its own words ("canceling statement due to statement timeout", which +a client cancel never produces — that one reads "due to user request"). + +Opt-in: it **skips unless `AGAMI_IT_PG_PASSWORD` is set** (so normal CI, which has no Postgres, is +unaffected, and no test password lives in the source). To run it against the integration fixture: + + docker compose -f tests/integration/docker-compose.yml up -d postgres + AGAMI_IT_PG_PASSWORD= \ + uv run pytest tests/test_postgres_timeout_integration.py + +Host/port/user/db default to the fixture's values and can be overridden via the other AGAMI_IT_PG_* +vars, exactly as in `test_postgres_named_cursor_integration.py`. +""" + +from __future__ import annotations + +import contextlib +import os +import sys +import threading +import time +import urllib.parse +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +PKG_SRC = REPO_ROOT / "packages" / "agami-core" / "src" +if str(PKG_SRC) not in sys.path: + sys.path.insert(0, str(PKG_SRC)) + +import execute_sql # noqa: E402 +import guardrail # noqa: E402 + +PROFILE = "analytics" + +# Long enough that the observer below reliably catches the statement mid-flight on a loaded machine, +# short enough that the whole file stays a few seconds. The resolver deals in whole seconds, so this +# is also the smallest useful step above the 1s floor the SQLite tests use. +_BUDGET_S = 3 + +# A statement with no termination in reach on any machine: a three-way cross join of ten thousand +# rows each, so a trillion nested-loop iterations of pure CPU. Sized that way on purpose — "bounded +# rather than hanging" is only proved if the unbounded run would take far longer than the assertion +# allows. +# +# The SHAPE matters as much as the size. One wide `generate_series(1, 100000000000)` looks like the +# simpler runaway and is wrong: Postgres materializes a function scan into a tuplestore, so that +# statement burns time by filling the server's temp tablespace, and a test that leaves the disk full +# has done more damage than the bug it was checking for. Three narrow series each materialize a few +# hundred kilobytes and never grow — the cost is entirely time, which is the only axis under test. +# +# `pg_sleep` would be the obvious runaway and is deliberately not used: the read-only guard blocks it +# by name (it is a DoS function), so it would never reach the executor. This does reach it — it opens +# with SELECT, names no physical table, projects no star, and calls nothing the guard denies. +_RUNAWAY_SQL = ( + "SELECT count(*) AS n FROM generate_series(1, 10000) AS a(x), " + "generate_series(1, 10000) AS b(y), generate_series(1, 10000) AS c(z)" +) + +# What the executor's work looks like from the server side. The Postgres path runs through a +# server-side cursor, so the backend's `query` is the `DECLARE "agami_bounded" …` and then the +# `FETCH FORWARD … FROM "agami_bounded"` that actually burns the CPU — the cursor name, not the +# SELECT text, is what is on the wire for most of the run. +_CURSOR_NAME = "agami_bounded" + + +def _pg_creds() -> dict[str, str]: + # The password is env-only (no source-embedded test secret); host/port/user/db default to the + # fixture's values. Override any of them via AGAMI_IT_PG_* to point at another Postgres. + return { + "type": "postgres", + "host": os.environ.get("AGAMI_IT_PG_HOST", "127.0.0.1"), + "port": os.environ.get("AGAMI_IT_PG_PORT", "55432"), + "user": os.environ.get("AGAMI_IT_PG_USER", "agami_test"), + "password": os.environ["AGAMI_IT_PG_PASSWORD"], + "database": os.environ.get("AGAMI_IT_PG_DB", "shop"), + } + + +@pytest.fixture +def pg_observer(): + """A live Postgres, plus the SECOND connection the assertions are made from. + + `autocommit` is not a detail: `pg_stat_activity` is served from a per-transaction snapshot of the + backend-status array, so polling it inside one long transaction returns the same stale rows + forever. Each poll has to be its own transaction to see the server's current state. + """ + psycopg2 = pytest.importorskip("psycopg2") + if not os.environ.get("AGAMI_IT_PG_PASSWORD"): + pytest.skip("set AGAMI_IT_PG_PASSWORD to run the live-Postgres integration test") + creds = _pg_creds() + try: + conn = psycopg2.connect( + host=creds["host"], port=int(creds["port"]), user=creds["user"], + password=creds["password"], dbname=creds["database"], connect_timeout=3, + ) + except Exception as exc: # no DB in this environment → skip, don't fail + pytest.skip(f"no reachable Postgres for the integration test ({exc})") + conn.autocommit = True + try: + yield conn, creds + finally: + conn.close() + + +@pytest.fixture(autouse=True) +def _reset_overrides(): + # Both overrides are request-scoped ContextVars; isolate this file from whatever set them. The + # budget below travels by env var rather than by ContextVar on purpose: the executor runs on a + # thread this test starts, and a new thread begins with an empty context. + execute_sql._timeout_override.set(None) + execute_sql._max_rows_override.set(None) + yield + execute_sql._timeout_override.set(None) + execute_sql._max_rows_override.set(None) + + +def _dsn(creds: dict[str, str]) -> str: + quote = urllib.parse.quote + return ( + f"postgresql://{quote(creds['user'])}:{quote(creds['password'])}" + f"@{creds['host']}:{creds['port']}/{quote(creds['database'])}" + ) + + +def _our_backends(conn, creds: dict[str, str]) -> list[tuple]: + """Every backend OTHER than this one whose work is the executor's — the server's own answer to + "is that statement still there?". `pg_backend_pid()` excludes the observer, whose own query text + contains the cursor name it is searching for.""" + with conn.cursor() as cur: + cur.execute( + "SELECT pid, state, left(query, 120) FROM pg_stat_activity " + "WHERE datname = %s AND pid <> pg_backend_pid() AND query LIKE %s", + (creds["database"], f"%{_CURSOR_NAME}%"), + ) + return cur.fetchall() + + +def _wait_until(predicate, limit_s: float, poll_s: float = 0.02) -> bool: + deadline = time.monotonic() + limit_s + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(poll_s) + return predicate() + + +def _guarded_on_a_thread(sql: str) -> tuple[threading.Thread, dict]: + outcome: dict = {} + + def call() -> None: + try: + outcome["env"] = execute_sql.execute_guarded( + sql, PROFILE, None, executor=execute_sql.BUILTIN_EXECUTOR, no_safety=True + ) + except BaseException as exc: # the chokepoint is total; record a breach rather than hide it + outcome["raised"] = exc + + worker = threading.Thread(target=call, name="agami-it-guarded", daemon=True) + worker.start() + return worker, outcome + + +def test_the_cancel_reaches_the_server_and_the_statement_is_gone(pg_observer, monkeypatch): + """The headline, and the one criterion the client cannot self-certify. + + Three facts in order, from a connection the executor does not own: the statement WAS running on + the server, the caller got the `resource_limit` refusal, and the statement is no longer there. + The last one is polled rather than asserted instantaneously — a cancel is a request the backend + acts on at its next check, not a synchronous return — but the window is two seconds, which a + statement counting to a hundred billion would still be inside by many minutes. + + The detail assertion is not decoration. `execute_guarded` has TWO time bounds, and the outer one + (`_execute_bounded`) is exactly the client-only abandon this test exists to rule out: it stops + waiting, leaves the worker inside the driver, and cannot say the statement stopped. Its refusal + says so in different words. Asserting the INNER wording is what proves the cancel is the thing + that ended this, rather than the caller walking away. + """ + conn, creds = pg_observer + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", _dsn(creds)) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + started = time.monotonic() + worker, outcome = _guarded_on_a_thread(_RUNAWAY_SQL) + + # 1. The server really is running it — otherwise "it is gone" proves nothing. + assert _wait_until(lambda: bool(_our_backends(conn, creds)), limit_s=_BUDGET_S), ( + "the statement never appeared in pg_stat_activity; the test would prove nothing" + ) + + # 2. The caller gets the refusal, well inside the outer bound (budget + 10s). + worker.join(timeout=_BUDGET_S + 30) + assert not worker.is_alive(), "the executor never returned" + assert "raised" not in outcome, outcome.get("raised") + elapsed = time.monotonic() - started + env = outcome["env"] + assert env.status == "refused", env + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert env.data is None + assert f"{_BUDGET_S}s limit" in env.refusal.detail, env.refusal.detail + assert "cancelled" in env.refusal.detail, env.refusal.detail # the inner bound, not the abandon + assert elapsed < _BUDGET_S + execute_sql._OUTER_BOUND_SKEW_S, f"{elapsed:.1f}s" + + # 3. The server's view: the statement is gone, not merely disowned. A backend left `idle in + # transaction` would still be listed here and would fail this assertion, so this covers the + # half-abandon (client stopped reading, session still holding the transaction open) as well + # as the outright runaway. + assert _wait_until(lambda: not _our_backends(conn, creds), limit_s=2.0), ( + f"still on the server after the refusal: {_our_backends(conn, creds)}" + ) + + +def test_the_native_statement_timeout_is_really_set_on_the_session(pg_observer, monkeypatch): + """`SET LOCAL statement_timeout` is the backstop for the case the watchdog cannot cover — our + process dying mid-statement — so it has to be proven to have LANDED, not just to have been sent. + + It cannot be read back the ordinary way: `SHOW statement_timeout` is not a SELECT and + `current_setting()` is a denied function, so neither reaches the executor, and the setting is + transaction-scoped so it is gone by the time anything else could look. What is left is to let it + fire. With our watchdog disarmed the statement runs on to the native bound, and the proof is + Postgres's own wording: it says "statement timeout", where a client `conn.cancel()` says "user + request". Only the server can produce the former. + + The elapsed floor pins the skew as well as the fact: the native bound is the budget PLUS five + seconds, deliberately behind the watchdog so the watchdog always wins the race in normal + operation and stays the sole classification signal. + """ + conn, creds = pg_observer + budget_s = 1 + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", _dsn(creds)) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(budget_s)) + + @contextlib.contextmanager + def _disarmed(cancel, timeout_s): + # The watchdog, minus the timer: the Event never fires, so nothing client-side stops this. + yield threading.Event() + + monkeypatch.setattr(execute_sql, "_deadline", _disarmed) + + started = time.monotonic() + worker, outcome = _guarded_on_a_thread(_RUNAWAY_SQL) + worker.join(timeout=budget_s + execute_sql._OUTER_BOUND_SKEW_S + 30) + assert not worker.is_alive(), "the executor never returned" + elapsed = time.monotonic() - started + + env = outcome["env"] + # A server-side kill is not our refusal: nothing set the flag, so it is an ordinary database + # failure — which is exactly why the watchdog is set to fire first in normal operation. + assert env.status == "failed", env + assert "statement timeout" in env.failure.message.lower(), env.failure.message + native_s = budget_s + execute_sql._NATIVE_BOUND_SKEW_S + assert elapsed >= native_s - 0.5, f"killed after {elapsed:.1f}s, before the {native_s}s bound" + assert _wait_until(lambda: not _our_backends(conn, creds), limit_s=2.0), ( + f"still on the server after the native timeout: {_our_backends(conn, creds)}" + ) From 82389065b914ef59171968309dab646a141e65d5 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 17:40:02 -0700 Subject: [PATCH 06/10] =?UTF-8?q?fix(executor):=20eight=20review=20finding?= =?UTF-8?q?s=20=E2=80=94=20the=20late=20cancel,=20the=20finally=20that=20e?= =?UTF-8?q?ats=20the=20refusal,=20the=20uncapped=20leak?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-statement deadline held on the happy path and leaked its guarantee on three edges either side of it. `_deadline`'s disarm was a request, not a join: `Timer.cancel()` only sets the timer's internal Event, so a `fire` that had already passed its own check ran anyway — setting the flag after the block returned, and delivering a cancel to a connection the engine had moved on from. A lock plus a disarmed flag make the disarm wait for a cancel in flight and make a losing `fire` land nowhere, so "checked after the watchdog is disarmed, so the flag is final" is now true. Three engines closed the connection unguarded in `finally`, where a raised exception REPLACES the one propagating — destroying the marker the `except _ResourceLimit: raise` above had just re-raised, and handing the caller an unclassified failure instead of the refusal. MySQL was the most exposed: its cancel destroys the socket that close then writes to. All ten now match. The abandoned worker the outer bound costs had no ceiling. On the injected path there is no inner watchdog, so every slow statement abandons one, and bounding the caller's wait removed the backpressure that used to cap them — leaks accumulated until the pool belonged to work nobody was waiting for. A process-wide cap refuses fast at the chokepoint instead of starting leak N+1, and says the executor is saturated rather than blaming a statement that never ran. `_timeout_override` is deleted rather than documented. It outranked the environment in the parent and could not cross the fork, so the supervisor bound derived from it could sit below the child's actual budget and fire first, inverting the ordered family. The budget now has exactly one configuration surface, pinned structurally and across a real subprocess. Also: the guardrail contract no longer understates the reach of the rule (all ten engines are wired; BigQuery's server-side-only bound is the one real residual); the resolver uses `isdecimal` so a superscript digit cannot raise out of an unguarded call site, and warns whenever the budget differs from what the operator wrote rather than only when the text was unreadable; the CLI adapter renders the marker as a refusal instead of a traceback. Tests: the fetch-inside-the-deadline criterion is now pinned on all nine watchdog engines rather than two (mutation-verified against sqlserver, oracle and snowflake); the both-surfaces test asserts an audit row per transport; the cancel-that-raises test polls instead of sleeping; and the short budget in the no-enumeration suite is scoped to the vector that needs it. --- packages/agami-core/src/execute_sql.py | 267 ++++++++--- packages/agami-core/src/guardrail.py | 6 +- packages/agami-core/src/mcp_http.py | 7 +- packages/agami-core/src/tools.py | 12 +- plugins/agami/lib/execute_sql.py | 267 ++++++++--- plugins/agami/lib/guardrail.py | 6 +- tests/test_ace035_no_enumeration.py | 21 +- tests/test_ace038_timeout.py | 494 +++++++++++++++++---- tests/test_postgres_timeout_integration.py | 8 +- 9 files changed, 854 insertions(+), 234 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index bd267e8..1b2f7b1 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -542,7 +542,14 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: cur.close() except Exception: pass - conn.close() + # And the connection close is guarded for the same reason, one line down: an exception raised + # inside a `finally` REPLACES the one propagating through it, so an unguarded close here would + # destroy the marker the `except _ResourceLimit: raise` above just re-raised — and the caller + # would read a bound we imposed as an unclassified server break. + try: + conn.close() + except Exception: + pass return result @@ -588,7 +595,15 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: - conn.close() + # Guarded because an exception raised inside a `finally` REPLACES the one propagating through + # it, which would destroy the marker re-raised just above. MySQL is the most exposed of the + # ten: its cancel is `_force_close()`, which deliberately destroys the socket that this + # `close()` then tries to write COM_QUIT to — so on exactly the timeout path this is the close + # most likely to raise. + try: + conn.close() + except Exception: + pass return result @@ -827,7 +842,13 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: - conn.close() + # Guarded like every other engine's: an exception raised inside a `finally` REPLACES the one + # propagating through it, so an unguarded close would silently convert the refusal above into + # an unclassified failure. + try: + conn.close() + except Exception: + pass return result @@ -1102,11 +1123,20 @@ def _resolve_row_cap() -> int: # bounds a whole process (interpreter start, model load, credential resolution, connect) rather than # a statement. _SUPERVISOR_SKEW_S = 60 -# Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the -# HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the -# context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one -# thread) it behaves exactly as a module global would. -_timeout_override: ContextVar[int | None] = ContextVar("_timeout_override", default=None) +# How many worker threads `_execute_bounded` may have given up on at once, process-wide. That layer +# holds nothing it can cancel, so every expiry abandons a thread which still occupies a pooled +# connection and, server-side, a running statement. On the built-in path the inner watchdog fires ten +# seconds earlier, so reaching the outer bound means a cancel already failed — rare, and an honest +# cost. On an INJECTED executor there is no inner watchdog at all, so every slow statement abandons +# one, and because the caller's own thread is freed at the bound the abandonments accumulate with no +# ceiling of their own: the anyio worker limiter used to supply that ceiling by blocking the caller, +# and bounding the wait removed it. The cap restores one explicitly. 8 is deliberately far below both +# that former limiter and a typical warehouse connection pool, so a saturated executor loses a +# minority of the pool to work nobody is waiting for rather than taking the datasource away from the +# whole organization. +_MAX_ABANDONED_WORKERS = 8 +_abandoned_lock = threading.Lock() +_abandoned_workers = 0 def _resolve_timeout_s() -> int: @@ -1115,25 +1145,32 @@ def _resolve_timeout_s() -> int: availability tradeoff and may set it higher OR lower than 30. A missing or non-positive value falls back to the default. - Unlike `_resolve_row_cap`, a value that is PRESENT but unparseable is logged at warning before - the fallback. An operator who wrote `45.5` or `30s` asked for something specific and silently - running 30 instead is how a misconfiguration survives a whole deployment unnoticed. The warning - goes to the module logger and never to stderr, because the subprocess transport parses stderr and - an extra line there would break that contract.""" + **The environment is the ONLY source, deliberately.** A request-scoped override would outrank it + in the parent and be invisible to a forked child, which re-resolves from `os.environ` alone — so + the supervisor bound the parent derives could sit BELOW the budget the child actually enforces + and fire first, inverting the ordered family the whole design rests on. One source, readable on + both sides of the fork, makes that inversion unrepresentable rather than merely unlikely. + + Unlike `_resolve_row_cap`, a value that is PRESENT and does not survive to become the budget is + logged at warning before the fallback. That covers `45.5` and `30s`, which cannot be read at all, + and equally `-5` and `0`, which can be read and are then declined: an operator who wrote either + asked for something specific, and a deployment quietly running 30 instead is exactly the + invisible degradation the warning exists against. The warning goes to the module logger and never + to stderr, because the subprocess transport parses stderr and an extra line there would break + that contract.""" raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() digits = raw[1:] if raw.startswith("-") else raw # a leading minus is a value, not a typo - if raw and not digits.isdigit(): + # `isdecimal`, not `isdigit`: the latter admits `²` and `①`, which `int()` then refuses — turning + # a misconfigured deployment into a ValueError raised out of this resolver, at a call site (the + # fork path's supervisor bound) that sits outside any handler. + written = int(raw) if digits.isdecimal() else None + timeout_s = written if written is not None and written > 0 else _DEFAULT_TIMEOUT_S + if raw and timeout_s != written: _LOG.warning( - "AGAMI_SQL_TIMEOUT_S=%r is not a whole number of seconds; falling back to %ds.", + "AGAMI_SQL_TIMEOUT_S=%r is not a usable whole number of seconds; falling back to %ds.", raw, _DEFAULT_TIMEOUT_S, ) - timeout_s = int(raw) if digits.isdigit() else _DEFAULT_TIMEOUT_S - if timeout_s <= 0: - timeout_s = _DEFAULT_TIMEOUT_S # "0" / "-5" → the default, never an instantly-expired budget - override = _timeout_override.get() - if override is not None and override > 0: - timeout_s = override # a caller that resolved its own budget outranks the deployment default return timeout_s @@ -1155,6 +1192,16 @@ class _OuterBoundExpired(_ResourceLimit): """ +class _ExecutorSaturated(_ResourceLimit): + """Raised INSTEAD of starting a statement, when the abandoned-worker cap is already reached. + + A sibling of `_OuterBoundExpired` under the same parent for the same reason — one handler, one + refusal — and distinguishable for the same reason too: this statement did not run long, it did + not run at all. What we observed is the executor, not the statement, so the sentence the caller + reads must not blame the query it just sent. + """ + + # The marker's message, single-sourced because every engine raises it. It is diagnostic text, not # caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own # detail, so nothing a caller reads comes from here. @@ -1162,6 +1209,52 @@ class _OuterBoundExpired(_ResourceLimit): # The outer marker's message, and diagnostic in the same way — it names the executor rather than the # statement, because at this layer the statement is not the thing we observed. _OUTLIVED_OUTER_BOUND = "the executor outlived the outer bound around it" +# The saturation marker's message, diagnostic in the same way again. +_EXECUTOR_SATURATED = "the executor already has its limit of abandoned calls outstanding" + +# The remediation names only what would make THIS statement executable: on the served path the caller +# is an assistant with no shell and no deployment, so naming the environment variable would be advice +# it cannot take, addressed to someone who is not reading. +_NARROW_IT = ("Narrow the time range, reduce the grouping, or add a selective filter, " + "then run it again.") + + +def _resource_limit_refusal(exc: _ResourceLimit) -> Refusal: + """The ONE `resource_limit` refusal, built from whichever of the three bounds raised. + + Single-sourced because two entries into the engines lead here — the guarded chokepoint and the + subprocess/CLI adapter — and a caller must not be able to tell which one it came through. + + The budget is re-resolved rather than carried on the marker: the marker stays a plain exception, + and nothing between the engine call and here can change the environment the resolver reads, so it + is the same number the watchdog used. The configured number belongs in the detail — it is a + deployment setting, not a data value, and a bound the caller cannot see is one it cannot plan + around. + """ + timeout_s = _resolve_timeout_s() + if isinstance(exc, _ExecutorSaturated): + # Nothing ran, so nothing about THIS statement is the finding. Saying "your query was too + # slow" here would be false and would send the caller off simplifying a statement that is + # very possibly fine. + detail = ( + "The executor is saturated: too many earlier calls have not returned, so this " + "statement was not started." + ) + remediation = "Wait for the calls already in flight to finish, then run it again." + elif isinstance(exc, _OuterBoundExpired): + # The outer layer holds no connection, so it stopped WAITING rather than stopping the + # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the + # driver call and the statement may still be running. The caller is told what we actually + # know, and the number quoted is the bound that actually elapsed. + detail = ( + f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " + "and the query was abandoned." + ) + remediation = _NARROW_IT + else: + detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." + remediation = _NARROW_IT + return refuse(RULE_RESOURCE_LIMIT, detail=detail, remediation=remediation) @contextlib.contextmanager @@ -1173,15 +1266,32 @@ def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threadin cancellation provokes must be able to read an already-set flag and attribute the failure to us rather than to the database. A `cancel` that raises is swallowed and logged, because some drivers raise when cancelled from a thread other than the one running the statement, and an exception - escaping a timer thread is both unhandleable by the caller and invisible in the result.""" + escaping a timer thread is both unhandleable by the caller and invisible in the result. + + **The disarm is a JOIN, not a request.** `threading.Timer.cancel()` only sets the timer's internal + `finished` Event, so a `fire` that already passed its own check runs to completion regardless — + setting the flag and delivering a cancel AFTER this block has returned. Both consequences are + real: every engine re-reads the flag once the block has exited and would read a stale `False`, and + a cancel arriving late lands on a connection the engine has moved on from, which on a pooled one + is by then someone else's statement. The lock plus the flag close that: a `fire` that loses the + race neither sets the Event nor calls `cancel`, and a `fire` that wins holds the lock across the + cancel, so the disarm waits for it. The wait is bounded by the layer outside this one — a cancel + that hangs is what the outer bound around the whole executor call exists to survive.""" fired = threading.Event() + lock = threading.Lock() + # Assigned only below, in the enclosing scope, and only read inside `fire` — so the closure sees + # the current value with no `nonlocal` and no mutable box. + disarmed = False def fire() -> None: - fired.set() - try: - cancel() - except Exception as exc: - _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) + with lock: + if disarmed: + return # the block already returned; there is nothing of ours left to stop + fired.set() + try: + cancel() + except Exception as exc: + _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) timer = threading.Timer(timeout_s, fire) timer.daemon = True # a hung cancel must never hold the interpreter open at shutdown @@ -1190,6 +1300,8 @@ def fire() -> None: yield fired finally: timer.cancel() # a block that finished on time disarms the watchdog before it can fire + with lock: + disarmed = True def _flag_truncated(cap: int) -> None: @@ -1301,7 +1413,8 @@ def _resolve_guard_model(profile: str): def _write_refusal(refusal: Refusal) -> None: """Write a guard refusal to stderr in the ONE shape every caller parses — the single wire-writer, - called only from ``main``. + called from ``main`` and from the per-engine CSV adapter, which is the other entry into the + engines. One JSON object, on one line: ``{"refusal": {reason, rule, detail, remediation}}`` — the wire shape S2 established, which ``tools._stderr_refusal`` rebuilds through ``Refusal`` on the parent @@ -1447,9 +1560,18 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa def _emit_or_err(run: Callable[[], ExecResult]) -> int: """Subprocess/CLI adapter over a ``_run_`` function: write its result to stdout as CSV and return exit code 0, or translate an ``ExecutorError`` into the stderr message + exit code the CLI - contract documents (byte-identical to what the old ``_execute_`` emitted).""" + contract documents (byte-identical to what the old ``_execute_`` emitted). + + The watchdog marker gets an arm of its own, AHEAD of the classified one, because this is the + second entry into the engine functions and they raise it. Without it a per-statement timeout + reached through here would escape as a traceback rather than as the exit code this docstring + promises. It renders exactly what ``main`` renders for a refusal — the one JSON object on stderr + and exit 1 — so the two entries cannot disagree about what a bound we imposed looks like.""" try: _emit_result_csv(run()) + except _ResourceLimit as exc: + _write_refusal(_resource_limit_refusal(exc)) + return 1 except ExecutorError as e: return _err(e.msg, code=e.code) return 0 @@ -1598,12 +1720,24 @@ def _execute_bounded( arbitrary ``Executor`` exposes one method and no connection, so the only thing this code owns is its own WAIT — it starts the call on a daemon thread and joins with the budget. - **The cost is a leaked worker, and it is real.** On expiry the thread is still inside the + **The cost is an abandoned worker, and it is real.** On expiry the thread is still inside the driver call, and it stays there: nothing here cancels it, and it may hold a database connection (and, on the server side, a running statement) until that call returns on its own. It is a - daemon so it cannot hold the interpreter open at exit, and that is the whole of the mitigation. - This is a bound on how long a CALLER waits, not a promise that the work stopped — which is - exactly why the inner watchdog, the layer that really can cancel, is set to fire first. + daemon so it cannot hold the interpreter open at exit. This is a bound on how long a CALLER + waits, not a promise that the work stopped — which is exactly why the inner watchdog, the layer + that really can cancel, is set to fire first. + + **And the cost is CAPPED, because bounding the wait removed the ceiling that used to bound it.** + Before this layer existed the same slow call blocked the caller's own worker thread, so the + host's thread limiter capped how much abandoned work could be in flight and applied backpressure + to everything behind it. Returning at the bound frees that slot, so on the injected path — where + there is no inner watchdog and EVERY slow statement abandons one — the abandonments would + otherwise accumulate without limit until the pool, and with it the datasource, belonged entirely + to work nobody is waiting for. So the count is capped by ``_MAX_ABANDONED_WORKERS`` and checked + BEFORE a worker is started: at the cap this refuses immediately rather than starting abandonment + N+1, and says the executor is saturated rather than blaming a statement that never ran. The slot + is released when the abandoned worker finally returns, which is the moment the leak actually + ends. Exceptions cross the thread boundary with their ORIGINAL type, re-raised here. That is what keeps the handlers in ``execute_guarded`` correct: ``_ResourceLimit`` still reaches the refusal @@ -1614,24 +1748,46 @@ def _execute_bounded( turn that fail-closed answer into a silent hang until the bound expired. The call runs inside a copy of the CALLER's context. A new thread starts with an empty one, so - without this the request-scoped ``_timeout_override`` / ``_max_rows_override`` would read as - unset inside the worker and every in-process call would silently fall back to the deployment - defaults. + without this the request-scoped ``_max_rows_override`` would read as unset inside the worker and + every in-process call would silently fall back to the deployment default. """ + global _abandoned_workers + timeout_s = _resolve_timeout_s() outcome: dict[str, Any] = {} ctx = copy_context() + with _abandoned_lock: + if _abandoned_workers >= _MAX_ABANDONED_WORKERS: + # Fail closed, and before the work starts: the cheapest moment to say no, and the only + # one at which saying no still prevents anything. + raise _ExecutorSaturated(_EXECUTOR_SATURATED) + def call() -> None: + global _abandoned_workers try: outcome["result"] = ctx.run(executor.execute, sql, creds, profile=profile) except BaseException as exc: outcome["error"] = exc + finally: + # `finished` and the slot release are set under the same lock the abandonment takes, so + # the two cannot both claim this call: a worker that finishes in the sliver between the + # join timing out and the lock being taken is counted as having returned, and its result + # is used rather than a refusal invented over the top of it. + with _abandoned_lock: + if outcome.get("abandoned"): + _abandoned_workers -= 1 + outcome["finished"] = True worker = threading.Thread(target=call, name="agami-bounded-execute", daemon=True) worker.start() worker.join(timeout_s + _OUTER_BOUND_SKEW_S) - if worker.is_alive(): + with _abandoned_lock: + abandoned = not outcome.get("finished") + if abandoned: + _abandoned_workers += 1 + outcome["abandoned"] = True + if abandoned: raise _OuterBoundExpired(_OUTLIVED_OUTER_BOUND) if "error" in outcome: raise outcome["error"] @@ -1726,38 +1882,11 @@ def execute_guarded( # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a # fix we can honestly name, unlike the supervisor's kill of a child that never returned. # - # ONE handler for both bounds, and so exactly one refusal per call however many layers were - # armed. When the inner watchdog fired, its marker is what arrives here — the outer layer's - # join returned long before its own budget and has nothing to add, so the inner refusal is - # the answer and cannot be overwritten by a later one. - # - # The budget is re-resolved rather than carried on the marker: the marker stays a plain - # exception, and nothing between the engine call and here can change the env var or the - # request-scoped ContextVar the resolver reads, so it is the same number the watchdog used. - timeout_s = _resolve_timeout_s() - # The configured number belongs in the detail — it is a deployment setting, not a data - # value, and a bound the caller cannot see is one it cannot plan around. - if isinstance(exc, _OuterBoundExpired): - # The outer layer holds no connection, so it stopped WAITING rather than stopping the - # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the - # driver call and the statement may still be running. The caller is told what we - # actually know, and the number quoted is the bound that actually elapsed. - detail = ( - f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " - "and the query was abandoned." - ) - else: - detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." - return _envelope("refused", refusal=refuse( - RULE_RESOURCE_LIMIT, - detail=detail, - # The remediation names only what would make THIS statement executable: on the served - # path the caller is an assistant with no shell and no deployment, so naming the - # environment variable would be advice it cannot take, addressed to someone who is not - # reading. - remediation="Narrow the time range, reduce the grouping, or add a selective filter, " - "then run it again.", - )) + # ONE handler for all three bounds, and so exactly one refusal per call however many layers + # were armed. When the inner watchdog fired, its marker is what arrives here — the outer + # layer's join returned long before its own budget and has nothing to add, so the inner + # refusal is the answer and cannot be overwritten by a later one. + return _envelope("refused", refusal=_resource_limit_refusal(exc)) except ExecutorError as exc: # The classified branch: `msg` is authored by this module (a missing driver, a connect # failure, the credential-resolution remediation naming DATASOURCE_URL) or relayed from the diff --git a/packages/agami-core/src/guardrail.py b/packages/agami-core/src/guardrail.py index 8c0bc07..585d4ec 100644 --- a/packages/agami-core/src/guardrail.py +++ b/packages/agami-core/src/guardrail.py @@ -55,8 +55,10 @@ # subject is the statement, which is what earns it a rule at all — "narrow the query" is a fix we can # honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this rule: it # stops a child that never returned, without knowing what the child was doing when it stopped, so it -# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Wired into SQLite first; the -# other engines follow, and until they do a statement on them is bounded only by that supervisor. +# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Every engine the executor speaks +# to is wired, with one recorded residual: BigQuery has no connection to cancel, so there the bound is +# server-side only and a client-side stall comes back as the executor never returning rather than as a +# statement we stopped. RULE_RESOURCE_LIMIT = "resource_limit" RULE_UNPARSEABLE = "unparseable" diff --git a/packages/agami-core/src/mcp_http.py b/packages/agami-core/src/mcp_http.py index f032d14..1704ecb 100644 --- a/packages/agami-core/src/mcp_http.py +++ b/packages/agami-core/src/mcp_http.py @@ -428,9 +428,10 @@ async def _call_tool(name: str, arguments: dict) -> list: raised = False try: # Run the tool handler OFF the event loop. The heavy handlers block for the whole query — - # execute_sql shells out to a subprocess (up to 240 s) and hits the warehouse — so on the - # loop a single slow query would freeze every other in-flight request. This completes - # ACE-048 (which off-loaded the KDF/OIDC/audit calls but left the handler on the loop). + # execute_sql runs it under the executor's own bounds (the per-statement budget, plus the + # supervisor's slack on the fork path) and hits the warehouse — so on the loop a single + # slow query would freeze every other in-flight request. This completes ACE-048 (which + # off-loaded the KDF/OIDC/audit calls but left the handler on the loop). # `run_blocking` (anyio.to_thread) copies the request context into the worker thread, so # the org-scoped model cache (ACE-045, read via `_current_org_ctx`) stays tenant-correct. result_text = await run_blocking(meta["handler"], arguments or {}) diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index 4b92f1e..a7f9745 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -1419,8 +1419,16 @@ def tool_execute_sql(args: dict[str, Any]) -> str: # inverted that order for any statement budget approaching it: the supervisor fired FIRST, so a # statement we could have cancelled and refused precisely came back instead as a # `failed`/`timeout` naming nothing the caller can act on. Imported lazily for the same - # reason `_run_in_process` does it. The child inherits `os.environ` (no `env=` below), so it - # resolves the identical `AGAMI_SQL_TIMEOUT_S` this call just read. + # reason `_run_in_process` does it. + # + # Resolved HERE and enforced on a child that re-resolves for itself, which only works because the + # resolver reads the environment and nothing else: the child inherits `os.environ` (no `env=` + # below) and therefore reaches the identical number. A request-scoped override would be the one + # thing that could break that — it would outrank the environment on this side of the fork and be + # invisible on the other, so a parent bound of 65s could sit against a child budget of 300s and + # fire first, inverting the order this whole family exists to hold. There is deliberately no such + # override; `_resolve_timeout_s` documents why, and a test pins that the budget keeps exactly one + # configuration surface. import execute_sql supervisor_timeout_s = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index bd267e8..1b2f7b1 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -542,7 +542,14 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: cur.close() except Exception: pass - conn.close() + # And the connection close is guarded for the same reason, one line down: an exception raised + # inside a `finally` REPLACES the one propagating through it, so an unguarded close here would + # destroy the marker the `except _ResourceLimit: raise` above just re-raised — and the caller + # would read a bound we imposed as an unclassified server break. + try: + conn.close() + except Exception: + pass return result @@ -588,7 +595,15 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: - conn.close() + # Guarded because an exception raised inside a `finally` REPLACES the one propagating through + # it, which would destroy the marker re-raised just above. MySQL is the most exposed of the + # ten: its cancel is `_force_close()`, which deliberately destroys the socket that this + # `close()` then tries to write COM_QUIT to — so on exactly the timeout path this is the close + # most likely to raise. + try: + conn.close() + except Exception: + pass return result @@ -827,7 +842,13 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: - conn.close() + # Guarded like every other engine's: an exception raised inside a `finally` REPLACES the one + # propagating through it, so an unguarded close would silently convert the refusal above into + # an unclassified failure. + try: + conn.close() + except Exception: + pass return result @@ -1102,11 +1123,20 @@ def _resolve_row_cap() -> int: # bounds a whole process (interpreter start, model load, credential resolution, connect) rather than # a statement. _SUPERVISOR_SKEW_S = 60 -# Per-call timeout, the twin of `_max_rows_override` and request-scoped for the same reason: once the -# HTTP server runs execution in-process, concurrent handlers run in worker threads that each copy the -# context, so one request's budget can never stomp another's. In the subprocess/CLI (one process, one -# thread) it behaves exactly as a module global would. -_timeout_override: ContextVar[int | None] = ContextVar("_timeout_override", default=None) +# How many worker threads `_execute_bounded` may have given up on at once, process-wide. That layer +# holds nothing it can cancel, so every expiry abandons a thread which still occupies a pooled +# connection and, server-side, a running statement. On the built-in path the inner watchdog fires ten +# seconds earlier, so reaching the outer bound means a cancel already failed — rare, and an honest +# cost. On an INJECTED executor there is no inner watchdog at all, so every slow statement abandons +# one, and because the caller's own thread is freed at the bound the abandonments accumulate with no +# ceiling of their own: the anyio worker limiter used to supply that ceiling by blocking the caller, +# and bounding the wait removed it. The cap restores one explicitly. 8 is deliberately far below both +# that former limiter and a typical warehouse connection pool, so a saturated executor loses a +# minority of the pool to work nobody is waiting for rather than taking the datasource away from the +# whole organization. +_MAX_ABANDONED_WORKERS = 8 +_abandoned_lock = threading.Lock() +_abandoned_workers = 0 def _resolve_timeout_s() -> int: @@ -1115,25 +1145,32 @@ def _resolve_timeout_s() -> int: availability tradeoff and may set it higher OR lower than 30. A missing or non-positive value falls back to the default. - Unlike `_resolve_row_cap`, a value that is PRESENT but unparseable is logged at warning before - the fallback. An operator who wrote `45.5` or `30s` asked for something specific and silently - running 30 instead is how a misconfiguration survives a whole deployment unnoticed. The warning - goes to the module logger and never to stderr, because the subprocess transport parses stderr and - an extra line there would break that contract.""" + **The environment is the ONLY source, deliberately.** A request-scoped override would outrank it + in the parent and be invisible to a forked child, which re-resolves from `os.environ` alone — so + the supervisor bound the parent derives could sit BELOW the budget the child actually enforces + and fire first, inverting the ordered family the whole design rests on. One source, readable on + both sides of the fork, makes that inversion unrepresentable rather than merely unlikely. + + Unlike `_resolve_row_cap`, a value that is PRESENT and does not survive to become the budget is + logged at warning before the fallback. That covers `45.5` and `30s`, which cannot be read at all, + and equally `-5` and `0`, which can be read and are then declined: an operator who wrote either + asked for something specific, and a deployment quietly running 30 instead is exactly the + invisible degradation the warning exists against. The warning goes to the module logger and never + to stderr, because the subprocess transport parses stderr and an extra line there would break + that contract.""" raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() digits = raw[1:] if raw.startswith("-") else raw # a leading minus is a value, not a typo - if raw and not digits.isdigit(): + # `isdecimal`, not `isdigit`: the latter admits `²` and `①`, which `int()` then refuses — turning + # a misconfigured deployment into a ValueError raised out of this resolver, at a call site (the + # fork path's supervisor bound) that sits outside any handler. + written = int(raw) if digits.isdecimal() else None + timeout_s = written if written is not None and written > 0 else _DEFAULT_TIMEOUT_S + if raw and timeout_s != written: _LOG.warning( - "AGAMI_SQL_TIMEOUT_S=%r is not a whole number of seconds; falling back to %ds.", + "AGAMI_SQL_TIMEOUT_S=%r is not a usable whole number of seconds; falling back to %ds.", raw, _DEFAULT_TIMEOUT_S, ) - timeout_s = int(raw) if digits.isdigit() else _DEFAULT_TIMEOUT_S - if timeout_s <= 0: - timeout_s = _DEFAULT_TIMEOUT_S # "0" / "-5" → the default, never an instantly-expired budget - override = _timeout_override.get() - if override is not None and override > 0: - timeout_s = override # a caller that resolved its own budget outranks the deployment default return timeout_s @@ -1155,6 +1192,16 @@ class _OuterBoundExpired(_ResourceLimit): """ +class _ExecutorSaturated(_ResourceLimit): + """Raised INSTEAD of starting a statement, when the abandoned-worker cap is already reached. + + A sibling of `_OuterBoundExpired` under the same parent for the same reason — one handler, one + refusal — and distinguishable for the same reason too: this statement did not run long, it did + not run at all. What we observed is the executor, not the statement, so the sentence the caller + reads must not blame the query it just sent. + """ + + # The marker's message, single-sourced because every engine raises it. It is diagnostic text, not # caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own # detail, so nothing a caller reads comes from here. @@ -1162,6 +1209,52 @@ class _OuterBoundExpired(_ResourceLimit): # The outer marker's message, and diagnostic in the same way — it names the executor rather than the # statement, because at this layer the statement is not the thing we observed. _OUTLIVED_OUTER_BOUND = "the executor outlived the outer bound around it" +# The saturation marker's message, diagnostic in the same way again. +_EXECUTOR_SATURATED = "the executor already has its limit of abandoned calls outstanding" + +# The remediation names only what would make THIS statement executable: on the served path the caller +# is an assistant with no shell and no deployment, so naming the environment variable would be advice +# it cannot take, addressed to someone who is not reading. +_NARROW_IT = ("Narrow the time range, reduce the grouping, or add a selective filter, " + "then run it again.") + + +def _resource_limit_refusal(exc: _ResourceLimit) -> Refusal: + """The ONE `resource_limit` refusal, built from whichever of the three bounds raised. + + Single-sourced because two entries into the engines lead here — the guarded chokepoint and the + subprocess/CLI adapter — and a caller must not be able to tell which one it came through. + + The budget is re-resolved rather than carried on the marker: the marker stays a plain exception, + and nothing between the engine call and here can change the environment the resolver reads, so it + is the same number the watchdog used. The configured number belongs in the detail — it is a + deployment setting, not a data value, and a bound the caller cannot see is one it cannot plan + around. + """ + timeout_s = _resolve_timeout_s() + if isinstance(exc, _ExecutorSaturated): + # Nothing ran, so nothing about THIS statement is the finding. Saying "your query was too + # slow" here would be false and would send the caller off simplifying a statement that is + # very possibly fine. + detail = ( + "The executor is saturated: too many earlier calls have not returned, so this " + "statement was not started." + ) + remediation = "Wait for the calls already in flight to finish, then run it again." + elif isinstance(exc, _OuterBoundExpired): + # The outer layer holds no connection, so it stopped WAITING rather than stopping the + # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the + # driver call and the statement may still be running. The caller is told what we actually + # know, and the number quoted is the bound that actually elapsed. + detail = ( + f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " + "and the query was abandoned." + ) + remediation = _NARROW_IT + else: + detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." + remediation = _NARROW_IT + return refuse(RULE_RESOURCE_LIMIT, detail=detail, remediation=remediation) @contextlib.contextmanager @@ -1173,15 +1266,32 @@ def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threadin cancellation provokes must be able to read an already-set flag and attribute the failure to us rather than to the database. A `cancel` that raises is swallowed and logged, because some drivers raise when cancelled from a thread other than the one running the statement, and an exception - escaping a timer thread is both unhandleable by the caller and invisible in the result.""" + escaping a timer thread is both unhandleable by the caller and invisible in the result. + + **The disarm is a JOIN, not a request.** `threading.Timer.cancel()` only sets the timer's internal + `finished` Event, so a `fire` that already passed its own check runs to completion regardless — + setting the flag and delivering a cancel AFTER this block has returned. Both consequences are + real: every engine re-reads the flag once the block has exited and would read a stale `False`, and + a cancel arriving late lands on a connection the engine has moved on from, which on a pooled one + is by then someone else's statement. The lock plus the flag close that: a `fire` that loses the + race neither sets the Event nor calls `cancel`, and a `fire` that wins holds the lock across the + cancel, so the disarm waits for it. The wait is bounded by the layer outside this one — a cancel + that hangs is what the outer bound around the whole executor call exists to survive.""" fired = threading.Event() + lock = threading.Lock() + # Assigned only below, in the enclosing scope, and only read inside `fire` — so the closure sees + # the current value with no `nonlocal` and no mutable box. + disarmed = False def fire() -> None: - fired.set() - try: - cancel() - except Exception as exc: - _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) + with lock: + if disarmed: + return # the block already returned; there is nothing of ours left to stop + fired.set() + try: + cancel() + except Exception as exc: + _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) timer = threading.Timer(timeout_s, fire) timer.daemon = True # a hung cancel must never hold the interpreter open at shutdown @@ -1190,6 +1300,8 @@ def fire() -> None: yield fired finally: timer.cancel() # a block that finished on time disarms the watchdog before it can fire + with lock: + disarmed = True def _flag_truncated(cap: int) -> None: @@ -1301,7 +1413,8 @@ def _resolve_guard_model(profile: str): def _write_refusal(refusal: Refusal) -> None: """Write a guard refusal to stderr in the ONE shape every caller parses — the single wire-writer, - called only from ``main``. + called from ``main`` and from the per-engine CSV adapter, which is the other entry into the + engines. One JSON object, on one line: ``{"refusal": {reason, rule, detail, remediation}}`` — the wire shape S2 established, which ``tools._stderr_refusal`` rebuilds through ``Refusal`` on the parent @@ -1447,9 +1560,18 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa def _emit_or_err(run: Callable[[], ExecResult]) -> int: """Subprocess/CLI adapter over a ``_run_`` function: write its result to stdout as CSV and return exit code 0, or translate an ``ExecutorError`` into the stderr message + exit code the CLI - contract documents (byte-identical to what the old ``_execute_`` emitted).""" + contract documents (byte-identical to what the old ``_execute_`` emitted). + + The watchdog marker gets an arm of its own, AHEAD of the classified one, because this is the + second entry into the engine functions and they raise it. Without it a per-statement timeout + reached through here would escape as a traceback rather than as the exit code this docstring + promises. It renders exactly what ``main`` renders for a refusal — the one JSON object on stderr + and exit 1 — so the two entries cannot disagree about what a bound we imposed looks like.""" try: _emit_result_csv(run()) + except _ResourceLimit as exc: + _write_refusal(_resource_limit_refusal(exc)) + return 1 except ExecutorError as e: return _err(e.msg, code=e.code) return 0 @@ -1598,12 +1720,24 @@ def _execute_bounded( arbitrary ``Executor`` exposes one method and no connection, so the only thing this code owns is its own WAIT — it starts the call on a daemon thread and joins with the budget. - **The cost is a leaked worker, and it is real.** On expiry the thread is still inside the + **The cost is an abandoned worker, and it is real.** On expiry the thread is still inside the driver call, and it stays there: nothing here cancels it, and it may hold a database connection (and, on the server side, a running statement) until that call returns on its own. It is a - daemon so it cannot hold the interpreter open at exit, and that is the whole of the mitigation. - This is a bound on how long a CALLER waits, not a promise that the work stopped — which is - exactly why the inner watchdog, the layer that really can cancel, is set to fire first. + daemon so it cannot hold the interpreter open at exit. This is a bound on how long a CALLER + waits, not a promise that the work stopped — which is exactly why the inner watchdog, the layer + that really can cancel, is set to fire first. + + **And the cost is CAPPED, because bounding the wait removed the ceiling that used to bound it.** + Before this layer existed the same slow call blocked the caller's own worker thread, so the + host's thread limiter capped how much abandoned work could be in flight and applied backpressure + to everything behind it. Returning at the bound frees that slot, so on the injected path — where + there is no inner watchdog and EVERY slow statement abandons one — the abandonments would + otherwise accumulate without limit until the pool, and with it the datasource, belonged entirely + to work nobody is waiting for. So the count is capped by ``_MAX_ABANDONED_WORKERS`` and checked + BEFORE a worker is started: at the cap this refuses immediately rather than starting abandonment + N+1, and says the executor is saturated rather than blaming a statement that never ran. The slot + is released when the abandoned worker finally returns, which is the moment the leak actually + ends. Exceptions cross the thread boundary with their ORIGINAL type, re-raised here. That is what keeps the handlers in ``execute_guarded`` correct: ``_ResourceLimit`` still reaches the refusal @@ -1614,24 +1748,46 @@ def _execute_bounded( turn that fail-closed answer into a silent hang until the bound expired. The call runs inside a copy of the CALLER's context. A new thread starts with an empty one, so - without this the request-scoped ``_timeout_override`` / ``_max_rows_override`` would read as - unset inside the worker and every in-process call would silently fall back to the deployment - defaults. + without this the request-scoped ``_max_rows_override`` would read as unset inside the worker and + every in-process call would silently fall back to the deployment default. """ + global _abandoned_workers + timeout_s = _resolve_timeout_s() outcome: dict[str, Any] = {} ctx = copy_context() + with _abandoned_lock: + if _abandoned_workers >= _MAX_ABANDONED_WORKERS: + # Fail closed, and before the work starts: the cheapest moment to say no, and the only + # one at which saying no still prevents anything. + raise _ExecutorSaturated(_EXECUTOR_SATURATED) + def call() -> None: + global _abandoned_workers try: outcome["result"] = ctx.run(executor.execute, sql, creds, profile=profile) except BaseException as exc: outcome["error"] = exc + finally: + # `finished` and the slot release are set under the same lock the abandonment takes, so + # the two cannot both claim this call: a worker that finishes in the sliver between the + # join timing out and the lock being taken is counted as having returned, and its result + # is used rather than a refusal invented over the top of it. + with _abandoned_lock: + if outcome.get("abandoned"): + _abandoned_workers -= 1 + outcome["finished"] = True worker = threading.Thread(target=call, name="agami-bounded-execute", daemon=True) worker.start() worker.join(timeout_s + _OUTER_BOUND_SKEW_S) - if worker.is_alive(): + with _abandoned_lock: + abandoned = not outcome.get("finished") + if abandoned: + _abandoned_workers += 1 + outcome["abandoned"] = True + if abandoned: raise _OuterBoundExpired(_OUTLIVED_OUTER_BOUND) if "error" in outcome: raise outcome["error"] @@ -1726,38 +1882,11 @@ def execute_guarded( # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a # fix we can honestly name, unlike the supervisor's kill of a child that never returned. # - # ONE handler for both bounds, and so exactly one refusal per call however many layers were - # armed. When the inner watchdog fired, its marker is what arrives here — the outer layer's - # join returned long before its own budget and has nothing to add, so the inner refusal is - # the answer and cannot be overwritten by a later one. - # - # The budget is re-resolved rather than carried on the marker: the marker stays a plain - # exception, and nothing between the engine call and here can change the env var or the - # request-scoped ContextVar the resolver reads, so it is the same number the watchdog used. - timeout_s = _resolve_timeout_s() - # The configured number belongs in the detail — it is a deployment setting, not a data - # value, and a bound the caller cannot see is one it cannot plan around. - if isinstance(exc, _OuterBoundExpired): - # The outer layer holds no connection, so it stopped WAITING rather than stopping the - # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the - # driver call and the statement may still be running. The caller is told what we - # actually know, and the number quoted is the bound that actually elapsed. - detail = ( - f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " - "and the query was abandoned." - ) - else: - detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." - return _envelope("refused", refusal=refuse( - RULE_RESOURCE_LIMIT, - detail=detail, - # The remediation names only what would make THIS statement executable: on the served - # path the caller is an assistant with no shell and no deployment, so naming the - # environment variable would be advice it cannot take, addressed to someone who is not - # reading. - remediation="Narrow the time range, reduce the grouping, or add a selective filter, " - "then run it again.", - )) + # ONE handler for all three bounds, and so exactly one refusal per call however many layers + # were armed. When the inner watchdog fired, its marker is what arrives here — the outer + # layer's join returned long before its own budget and has nothing to add, so the inner + # refusal is the answer and cannot be overwritten by a later one. + return _envelope("refused", refusal=_resource_limit_refusal(exc)) except ExecutorError as exc: # The classified branch: `msg` is authored by this module (a missing driver, a connect # failure, the credential-resolution remediation naming DATASOURCE_URL) or relayed from the diff --git a/plugins/agami/lib/guardrail.py b/plugins/agami/lib/guardrail.py index 8c0bc07..585d4ec 100644 --- a/plugins/agami/lib/guardrail.py +++ b/plugins/agami/lib/guardrail.py @@ -55,8 +55,10 @@ # subject is the statement, which is what earns it a rule at all — "narrow the query" is a fix we can # honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this rule: it # stops a child that never returned, without knowing what the child was doing when it stopped, so it -# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Wired into SQLite first; the -# other engines follow, and until they do a statement on them is bounded only by that supervisor. +# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Every engine the executor speaks +# to is wired, with one recorded residual: BigQuery has no connection to cancel, so there the bound is +# server-side only and a client-side stall comes back as the executor never returning rather than as a +# statement we stopped. RULE_RESOURCE_LIMIT = "resource_limit" RULE_UNPARSEABLE = "unparseable" diff --git a/tests/test_ace035_no_enumeration.py b/tests/test_ace035_no_enumeration.py index 2a124d8..7c347c8 100644 --- a/tests/test_ace035_no_enumeration.py +++ b/tests/test_ace035_no_enumeration.py @@ -180,11 +180,11 @@ def declared(tmp_path, monkeypatch): monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(artifacts)) monkeypatch.setenv("DATASOURCE_URL__ACME", f"sqlite:///{warehouse}") - # The smallest per-statement budget the resolver accepts, for the `resource_limit` vector: it is - # the one statement here that reaches the executor and has to be stopped there. Every other - # statement in this file is refused before execution or rejected instantly by the database, so a - # short budget costs them nothing. - monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") + # Deliberately NOT a one-second budget for every test that takes this fixture. Only the + # `resource_limit` vector wants one, and it sets it for itself; the `failed` vector reaches the + # same warehouse under it, so a stall on a loaded runner would turn an expected database + # rejection into a timeout refusal and fail a test that has nothing to do with the clock. + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) # Local, not hosted: the disk model is the one the gates use. (`model_unavailable` needs the # hosted signal and gets its own fixture.) monkeypatch.delenv("AGAMI_DB_URL", raising=False) @@ -370,7 +370,9 @@ def _assert_echo_only(body: dict, sql: str) -> None: @pytest.mark.parametrize(("rule", "sql", "route"), _MATRIX, ids=_MATRIX_IDS) -def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal(declared, rule, sql, route): +def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal( + declared, rule, sql, route, monkeypatch +): """Six rules — five here, `model_unavailable` below — across both surfaces and both execution paths. @@ -380,6 +382,13 @@ def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal(declared, ru by construction; the child's own read-only refusal crossing the wire is covered by `tests/test_ace035_read_only_refusal.py::test_parent_reconstructs_the_child_refusal`.) """ + if rule == guardrail.RULE_RESOURCE_LIMIT: + # The smallest per-statement budget the resolver accepts, scoped to the one vector that needs + # it: this is the only statement here that passes every gate, reaches the executor, and has + # to be stopped there. The forked and stdio routes inherit `os.environ`, so setting it here + # reaches them too. + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") + body = ROUTES[route](sql) assert body["status"] == "refused", body diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index d6c8dab..9cad803 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -1,10 +1,12 @@ """Per-statement timeout — config resolution, the deadline primitive, and the refusal it produces. -`_resolve_timeout_s` answers "how long may one statement run", from `AGAMI_SQL_TIMEOUT_S`, the -`_timeout_override` ContextVar, and a 30s default; unlike the row cap it complains out loud when the -env value is present but unparseable. `_deadline` is the watchdog those seconds feed: it fires an -Event and calls a cancel callable when a block outlives its budget, and disarms cleanly when it does -not. +`_resolve_timeout_s` answers "how long may one statement run", from `AGAMI_SQL_TIMEOUT_S` and a 30s +default and from nothing else — one configuration surface, so the parent and the forked child reach +the same number. Unlike the row cap it complains out loud whenever the budget it returns is not the +one the operator wrote, whether the text was unreadable or merely declined. `_deadline` is the +watchdog those seconds feed: it fires an Event and calls a cancel callable when a block outlives its +budget, and disarms cleanly when it does not — where "cleanly" means the disarm is a JOIN, so a +watchdog that loses the race lands nowhere at all. The second half of this file proves the whole contract end to end on ONE engine — SQLite, chosen because it is in-process, needs no network, and `sqlite3.Connection.interrupt()` is a genuine cancel @@ -45,6 +47,7 @@ import threading import time import types +from contextvars import ContextVar from pathlib import Path from types import SimpleNamespace @@ -64,17 +67,30 @@ @pytest.fixture(autouse=True) -def _reset_override(): - # Both overrides are request-scoped ContextVars; isolate every test from them. The row cap is - # reset here too because the outer bound copies the caller's whole context into its worker, so a - # cap left set by one test is now visible to the next test's executor as well. - execute_sql._timeout_override.set(None) +def _reset_row_cap(): + # The row cap is a request-scoped ContextVar; isolate every test from it. It matters more than it + # used to because the outer bound copies the caller's whole context into its worker, so a cap + # left set by one test is now visible to the next test's executor as well. execute_sql._max_rows_override.set(None) yield - execute_sql._timeout_override.set(None) execute_sql._max_rows_override.set(None) +@pytest.fixture(autouse=True) +def _drain_abandoned_workers(): + """Wait for any worker a test abandoned to finish, so its slot is back before the next test runs. + + The counter is process-wide and released by the abandoned worker itself, which returns shortly + after the test releases it. Draining here rather than zeroing the counter keeps the accounting + honest: setting it to 0 while a worker was still running would let that worker decrement past + zero and hand a later test a cap it has not actually got. + """ + yield + deadline = time.monotonic() + 5 + while execute_sql._abandoned_workers and time.monotonic() < deadline: + time.sleep(0.01) + + @pytest.fixture(autouse=True) def _clear_env(monkeypatch): # The suite must not inherit an operator's real budget from the ambient environment. @@ -123,8 +139,11 @@ def test_a_non_positive_value_falls_back_to_the_default(monkeypatch, raw): assert execute_sql._resolve_timeout_s() == 30 -@pytest.mark.parametrize("raw", ["6O", "30s", "45.5", "thirty", "1e3"]) -def test_an_unparseable_value_falls_back_and_says_so(monkeypatch, caplog, raw): +# `²` and `①` are the cases `str.isdigit()` waves through and `int()` then refuses. They are not a +# curiosity: the resolver is called at the fork path's supervisor bound, outside any handler, so a +# ValueError raised out of it escapes the tool edge as a traceback rather than as a budget. +@pytest.mark.parametrize("raw", ["6O", "30s", "45.5", "thirty", "1e3", "²", "①"]) +def test_an_unreadable_value_falls_back_and_says_so(monkeypatch, caplog, raw): """The row cap falls back silently; this one must not. An operator who typed a capital O for a zero has to be able to find out why their budget is not what they configured.""" monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) @@ -138,39 +157,79 @@ def test_an_unparseable_value_falls_back_and_says_so(monkeypatch, caplog, raw): ) -@pytest.mark.parametrize("raw", ["", "45", "0", "-5"]) -def test_a_parseable_or_absent_value_stays_quiet(monkeypatch, caplog, raw): - """Only genuinely unparseable text warrants the warning. A deliberate `0` or `-5` is a value we - understood and declined, not a typo, and warning on it would train operators to ignore the log.""" +@pytest.mark.parametrize("raw", ["0", "00", "-5", "-0"]) +def test_a_value_we_read_and_declined_says_so_too(monkeypatch, caplog, raw): + """The warning is about the OUTCOME, not about the parse. + + `-5` reads perfectly and is then declined, so it used to slip past silently while `6O` warned — + and a deployment quietly running 30s when its operator wrote something else is precisely the + invisible degradation the warning exists against. Whether we could read the text is not the + question the operator has; whether they got what they asked for is. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + assert execute_sql._resolve_timeout_s() == 30 + + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert warnings, f"no warning emitted for the declined value {raw!r}" + assert any(raw in r.getMessage() for r in warnings), ( + f"the warning must name the declined text {raw!r}; got {[r.getMessage() for r in warnings]}" + ) + + +@pytest.mark.parametrize("raw", ["", "45", "045", " 45 ", "600"]) +def test_a_value_the_operator_actually_got_stays_quiet(monkeypatch, caplog, raw): + """The other direction, so the warning cannot become noise. A leading zero or surrounding + whitespace still yields exactly the number that was written, and warning on it would train + operators to ignore the log.""" monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): execute_sql._resolve_timeout_s() assert [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] == [] -def test_the_context_var_takes_precedence_over_the_env(monkeypatch): - monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") - execute_sql._timeout_override.set(7) - assert execute_sql._resolve_timeout_s() == 7 +def test_the_budget_has_exactly_one_configuration_surface(): + """The environment is the only source, and that is load-bearing rather than tidy. + A request-scoped override outranks the environment in THIS process and cannot cross a fork: the + child re-resolves from `os.environ` alone. So a parent that derived the supervisor's bound from + an override would compute a bound BELOW the budget the child actually enforces, the supervisor + would fire first, and the ordered family the whole design rests on would invert — turning a + precise refusal into a `failed`/`timeout` that names nothing. Asserted structurally, because the + hazard is the existence of a second surface, not any particular value in one. + """ + context_vars = { + name for name, value in vars(execute_sql).items() if isinstance(value, ContextVar) + } + assert context_vars == {"_max_rows_override"}, ( + "a second, higher-precedence configuration surface for the budget cannot cross the fork; " + f"found {sorted(context_vars)}" + ) -def test_the_context_var_takes_precedence_over_the_default(): - execute_sql._timeout_override.set(12) - assert execute_sql._resolve_timeout_s() == 12 +def test_the_supervisor_bound_exceeds_the_budget_a_real_child_resolves(monkeypatch): + """The inversion, driven across the actual process boundary rather than reasoned about. -def test_the_context_var_may_raise_the_budget_as_well_as_lower_it(monkeypatch): - """Unlike the row cap, which can only be tightened per call, the override wins outright.""" - monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "5") - execute_sql._timeout_override.set(90) - assert execute_sql._resolve_timeout_s() == 90 + The parent computes the supervisor's bound; a real forked interpreter, inheriting this + environment exactly as `subprocess.run` gives it one, resolves its own budget. The first must + exceed the second, or the outermost bound fires before the innermost. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "300") + parent_bound = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S + child = subprocess.run( + [sys.executable, "-c", "import execute_sql; print(execute_sql._resolve_timeout_s())"], + capture_output=True, text=True, timeout=60, + env={**os.environ, "PYTHONPATH": str(PKG_SRC)}, + ) + assert child.returncode == 0, child.stderr + child_budget = int(child.stdout.strip()) -@pytest.mark.parametrize("override", [0, -1]) -def test_a_non_positive_override_is_ignored(monkeypatch, override): - monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") - execute_sql._timeout_override.set(override) - assert execute_sql._resolve_timeout_s() == 45 + assert child_budget == 300, child.stdout + assert parent_bound > child_budget, ( + f"the supervisor stops waiting at {parent_bound}s while the child is still inside a " + f"{child_budget}s budget" + ) # -------------------------------------------------------------------------------------------- @@ -233,22 +292,115 @@ def test_the_timer_is_disarmed_even_when_the_block_raises(): assert cancel.calls == 0 +class _ManualTimer: + """A `threading.Timer` stand-in whose `cancel()` does nothing and whose callback the test fires. + + Not a caricature of the real one — a faithful model of its worst case. `Timer.cancel()` only sets + the timer's internal `finished` Event, so a timer thread that has ALREADY passed its own + `if not self.finished.is_set()` check runs the callback regardless. This makes that interleaving + — disarm first, `fire` second — happen on every run rather than once in a thousand. + """ + + def __init__(self, interval, function): + self.interval = interval + self.function = function + self.started = False + + def start(self) -> None: + self.started = True + + def cancel(self) -> None: + pass + + +def test_a_watchdog_that_loses_the_race_to_the_disarm_lands_nowhere(monkeypatch): + """The disarm has to be a JOIN, not a request, and `Timer.cancel()` is only ever a request. + + Two things go wrong when a `fire` that lost the race still runs. Every engine re-reads the flag + once the block has exited — "checked after the watchdog is disarmed, so the flag is final" — and + a flag set afterwards makes a completed statement look like a refusal. And the cancel itself + lands on a connection the engine has moved on from; on a pooled one that is somebody else's + statement being killed by our watchdog. + """ + timers: list = [] + + def _timer(interval, function): + timers.append(_ManualTimer(interval, function)) + return timers[-1] + + monkeypatch.setattr(execute_sql.threading, "Timer", _timer) + cancel = _RecordingCancel() + + with execute_sql._deadline(cancel, _TINY) as fired: + pass + + assert timers and timers[0].started, "no watchdog was armed" + timers[0].function() # the timer thread had already committed; `fire` runs anyway + + assert cancel.calls == 0, "a cancel landed after the block the deadline was bounding had exited" + assert not fired.is_set(), "the flag every engine re-reads after the block was not final" + + +def test_the_disarm_waits_for_a_cancel_already_in_flight(): + """The other half of the same guarantee: a `fire` that WON the race is finished with before the + block is allowed to return. + + Otherwise "no late cancel arrives" is only true of the cancels that had not started yet, and a + driver cancel that takes a moment still reaches into whatever the connection does next. + """ + entered = threading.Event() + release = threading.Event() + exited = threading.Event() + + def cancel() -> None: + entered.set() + release.wait(5) + + def run() -> None: + with execute_sql._deadline(cancel, _TINY): + assert entered.wait(5), "the watchdog never ran" + exited.set() + + runner = threading.Thread(target=run, daemon=True) + runner.start() + assert entered.wait(5), "the watchdog never ran" + + assert not exited.wait(0.3), "the block returned while its own cancel was still running" + release.set() + runner.join(5) + assert exited.is_set() + + +def _wait_for_warning(caplog, needle: str, timeout_s: float = 5.0) -> bool: + """Poll the captured records for a warning containing `needle`, up to a deadline.""" + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if any(needle in r.getMessage() + for r in caplog.records if r.levelno == logging.WARNING): + return True + time.sleep(0.005) + return False + + def test_a_cancel_that_raises_does_not_escape_the_timer_thread(caplog): """Some drivers raise when cancelled from a thread other than the one running the statement. That must be logged and swallowed: an exception escaping a timer thread is unhandleable by the caller - and would be lost to threading's excepthook.""" + and would be lost to threading's excepthook. + + Polled rather than slept on. `cancel.done` is set in a `finally` that runs BEFORE the raise + reaches `fire`'s handler, so nothing this test can wait on is synchronized with the log write, + and a fixed sleep is a wager on how the runner happened to schedule two threads. + """ cancel = _RecordingCancel(fails=True) with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): with execute_sql._deadline(cancel, _TINY) as fired: assert cancel.done.wait(2.0), "the watchdog never ran" - time.sleep(_TINY) # let `fire` finish handling the raised cancel before we leave + assert _wait_for_warning(caplog, "driver refused to cancel"), ( + f"the failed cancel was not logged; got {[r.getMessage() for r in caplog.records]}" + ) assert fired.is_set() # the timeout still counts as fired even though the cancel failed assert cancel.calls == 1 - assert any( - r.levelno == logging.WARNING and "driver refused to cancel" in r.getMessage() - for r in caplog.records - ), f"the failed cancel was not logged; got {[r.getMessage() for r in caplog.records]}" def test_the_watchdog_thread_is_a_daemon_and_does_not_hold_the_process_open(): @@ -325,14 +477,14 @@ def _guarded(sql: str) -> object: ) -def test_a_runaway_statement_is_cancelled_rather_than_left_to_run(warehouse): +def test_a_runaway_statement_is_cancelled_rather_than_left_to_run(warehouse, monkeypatch): """The headline: a statement that would run for minutes is stopped at its budget and comes back as a refusal naming the rule the contract reserves for a bound we imposed. The elapsed assertion is the one that would still fail if the deadline were never armed — without it a test that merely waited out the query would look identical and pass in several minutes. """ - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) started = time.monotonic() env = _guarded(_RUNAWAY_SQL) @@ -346,14 +498,14 @@ def test_a_runaway_statement_is_cancelled_rather_than_left_to_run(warehouse): assert env.refusal.reason == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] -def test_a_cancelled_statement_yields_no_partial_data(warehouse): +def test_a_cancelled_statement_yields_no_partial_data(warehouse, monkeypatch): """Whatever rows the engine had gathered when the watchdog fired are not an answer. A truncated result presented as a result is the failure mode the bounded-fetch work already guards against on the row axis; on the time axis the answer is stronger — there is no data at all, and `Envelope.__post_init__` enforces that a refusal cannot carry any. """ - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) env = _guarded(_RUNAWAY_SQL) @@ -371,14 +523,14 @@ def execute(self, vetted_sql: str, creds: dict, *, profile: str): raise execute_sql._ResourceLimit("the statement outlived its per-statement budget") -def test_the_detail_quotes_the_configured_budget(warehouse): +def test_the_detail_quotes_the_configured_budget(warehouse, monkeypatch): """A bound the caller cannot see is one it cannot plan around, so the number is in the detail. The configured value is not a data value: it is a deployment setting, and stating it discloses nothing about the database or its contents. Asserted against a distinctive budget rather than the default, so a hard-coded `30s` in the message cannot pass. """ - execute_sql._timeout_override.set(7) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(7)) env = execute_sql.execute_guarded( "SELECT id FROM orders", PROFILE, None, @@ -390,7 +542,7 @@ def test_the_detail_quotes_the_configured_budget(warehouse): assert "7s" in env.refusal.detail, env.refusal.detail -def test_the_remediation_names_no_deployment_environment_variable(warehouse): +def test_the_remediation_names_no_deployment_environment_variable(warehouse, monkeypatch): """The remediation has to be addressed to whoever is reading it. On the served path that is an assistant holding a statement, with no shell, no deployment and no @@ -398,7 +550,7 @@ def test_the_remediation_names_no_deployment_environment_variable(warehouse): caller at an operator who is not in the conversation, and it reads as a fix while being unfollowable. What is left has to be something that would make THIS statement executable. """ - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) env = execute_sql.execute_guarded( "SELECT id FROM orders", PROFILE, None, @@ -482,14 +634,16 @@ def _connect(path, *a, **kw): return _install -def test_a_slow_fetch_is_bounded_even_when_the_execute_returned_at_once(warehouse, fake_sqlite): +def test_a_slow_fetch_is_bounded_even_when_the_execute_returned_at_once( + warehouse, fake_sqlite, monkeypatch +): """The clock covers the whole statement, fetch included — an explicit criterion, not a bonus. Bounding only `execute` would leave the common streaming shape unbounded: the driver returns immediately and the engine scans while the caller pulls. The cancel has to land on the fetch, and the refusal has to be the same one a slow execute produces. """ - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) conn = fake_sqlite(_SlowFetchCursor) started = time.monotonic() @@ -528,14 +682,16 @@ def fetchmany(self, n: int): # pragma: no cover - execute raises first raise AssertionError("unreachable") -def test_a_database_error_with_the_flag_unset_is_a_failure_not_a_refusal(warehouse, fake_sqlite): +def test_a_database_error_with_the_flag_unset_is_a_failure_not_a_refusal( + warehouse, fake_sqlite, monkeypatch +): """Direction (a): the watchdog never fired, so this is the database's outcome, not ours. A generous budget means the flag stays clear while the identical error text arrives. It must unwind as an `ExecutorError` and leave the chokepoint as `failed`/`syntax` — a refusal here would tell a caller to narrow a statement that never ran long at all. """ - execute_sql._timeout_override.set(300) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(300)) fake_sqlite(_ImmediateErrorCursor) env = _guarded("SELECT c FROM orders") @@ -579,7 +735,7 @@ def test_an_error_that_merely_arrives_late_is_still_a_failure(warehouse, fake_sq `failed`, because we did not stop this statement; it stopped on its own, late. """ monkeypatch.setattr(execute_sql, "_deadline", _deadline_that_never_fires) - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) fake_sqlite(_LateErrorCursor) started = time.monotonic() @@ -651,7 +807,7 @@ def audited(tmp_path, monkeypatch): return SimpleNamespace(app_db=app_db) -def test_the_refusal_is_written_to_the_audit_trail(audited): +def test_the_refusal_is_written_to_the_audit_trail(audited, monkeypatch): """A decision we made against a caller's statement is exactly the row a reviewer comes looking for, so this outcome is audited like every other — and keyed by the id the caller was handed. @@ -662,7 +818,7 @@ def test_the_refusal_is_written_to_the_audit_trail(audited): from store import Store tools.set_injected_executor(execute_sql.BUILTIN_EXECUTOR) - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) try: body = json.loads(tools.tool_execute_sql({"sql": _RUNAWAY_SQL, "datasource": PROFILE, "raw_query": "how many"})) @@ -728,7 +884,7 @@ def test_an_injected_executor_that_never_returns_is_still_bounded(warehouse, mon which needs no clock at all. """ monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) executor = _BlockingExecutor() started = time.monotonic() @@ -754,7 +910,7 @@ def test_the_outer_refusal_says_what_actually_happened(warehouse, monkeypatch): The sentence the caller reads names what we observed — the executor did not come back — and quotes the bound that actually elapsed rather than the inner budget.""" monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) executor = _BlockingExecutor() try: @@ -772,7 +928,65 @@ def test_the_outer_refusal_says_what_actually_happened(warehouse, monkeypatch): assert "AGAMI_" not in authored, authored -def test_the_inner_refusal_is_the_one_the_caller_gets(warehouse): +def _guarded_with(executor) -> object: + return execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=executor, no_safety=True, + ) + + +def _drain(deadline_s: float = 5.0) -> None: + """Wait for the abandoned workers to return and give their slots back.""" + deadline = time.monotonic() + deadline_s + while execute_sql._abandoned_workers and time.monotonic() < deadline: + time.sleep(0.01) + + +def test_the_abandoned_workers_are_capped_and_the_cap_releases(warehouse, monkeypatch): + """The cost of the outer bound is an abandoned worker, and it has to have a ceiling. + + Before this layer existed, a call this slow blocked the host's own worker thread — so the thread + limiter capped how much abandoned work could exist and pushed back on everything behind it. + Returning at the bound frees that slot, and on the injected path (no inner watchdog, so EVERY + slow statement abandons one) the abandonments would accumulate at the rate callers arrive, each + still holding a pooled connection and a running statement, until the datasource belonged to work + nobody is waiting for. + + So at the cap the chokepoint refuses BEFORE starting one more — fast, fail-closed, and saying the + executor is saturated rather than blaming a statement that never ran. And the slot comes back + when the abandoned worker finally returns, which is the moment the leak actually ends. + """ + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + monkeypatch.setattr(execute_sql, "_MAX_ABANDONED_WORKERS", 2) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + executor = _BlockingExecutor() + + try: + for n in range(2): + env = _guarded_with(executor) + assert env.status == "refused", (n, env) + assert "did not return" in env.refusal.detail, (n, env.refusal.detail) + assert execute_sql._abandoned_workers == 2 + + started = time.monotonic() + env = _guarded_with(executor) + elapsed = time.monotonic() - started + + assert elapsed < _BUDGET_S, f"the capped call still waited out a bound ({elapsed:.2f}s)" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert "saturated" in env.refusal.detail, env.refusal.detail + # It must not blame a statement that never ran, and it must not claim a cancel. + assert "cancelled" not in env.refusal.detail, env.refusal.detail + assert executor.calls == 2, "the refused call started leak N+1 anyway" + finally: + executor.release.set() + + _drain() + assert execute_sql._abandoned_workers == 0, "the slots never came back" + assert _guarded_with(executor).status == "ok" # and the executor is usable again + + +def test_the_inner_refusal_is_the_one_the_caller_gets(warehouse, monkeypatch): """When both layers are armed the INNER one wins, and the outer neither overwrites its refusal nor mints a second. @@ -782,7 +996,7 @@ def test_the_inner_refusal_is_the_one_the_caller_gets(warehouse): armed and genuinely loses the race. The detail is the assertion, because it is the one part of the envelope the two layers word differently. """ - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) started = time.monotonic() env = _guarded(_RUNAWAY_SQL) @@ -866,7 +1080,7 @@ def test_a_plain_exception_still_reaches_the_catch_all_across_the_worker(warehou class _ContextProbeExecutor: - """Records the thread it ran on and what the two request-scoped ContextVars read there.""" + """Records the thread it ran on, the request-scoped row cap it read there, and the budget.""" def __init__(self) -> None: self.thread: threading.Thread | None = None @@ -875,21 +1089,24 @@ def __init__(self) -> None: def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: self.thread = threading.current_thread() - self.seen_timeout = execute_sql._timeout_override.get() + self.seen_timeout = execute_sql._resolve_timeout_s() self.seen_rows = execute_sql._max_rows_override.get() return execute_sql.ExecResult(columns=["c"], rows=[(1,)], truncated=False) -def test_the_request_scoped_overrides_are_visible_inside_the_worker(warehouse): +def test_the_request_scoped_row_cap_is_visible_inside_the_worker(warehouse, monkeypatch): """A new thread starts with an EMPTY context, so without an explicit copy the caller's per-call - budget and row cap would read as unset inside the worker and every in-process call would quietly - fall back to the deployment defaults — a bug with no symptom other than the wrong numbers. + row cap would read as unset inside the worker and every in-process call would quietly fall back + to the deployment default — a bug with no symptom other than the wrong numbers. + + The budget is asserted alongside it for the opposite reason: it rides the environment, which a + new thread inherits for free, and the two must agree inside the worker as well as outside it. The thread identity is asserted first, because it is what makes the rest of this test mean - anything: if the call ever stopped running off-thread, the ContextVar assertions would pass + anything: if the call ever stopped running off-thread, the ContextVar assertion would pass trivially and the copy could be deleted with the suite still green. """ - execute_sql._timeout_override.set(7) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "7") execute_sql._max_rows_override.set(5) probe = _ContextProbeExecutor() @@ -1032,14 +1249,29 @@ def _http_refusal(sql: str) -> dict: return json.loads(resp.json()["result"]["content"][0]["text"]) +def _audit_rows(app_db: str) -> list: + from store import Store + + store = Store.connect(app_db) + try: + return store.query("SELECT id, status, reason, rule FROM query_executions") + finally: + store.close() + + def test_both_surfaces_refuse_a_runaway_statement_identically(audited, monkeypatch): - """The criterion in one test: the limit applies the same on stdio and on HTTP, and no executor - escapes it. + """The criterion in one test: the limit applies the same on stdio and on HTTP, no executor + escapes it, and BOTH leave an audit row. The two transports run DIFFERENT execution paths — stdio forks a child and reads its stderr back, HTTP runs the built-in executor in-process behind the outer bound — so "identically" is asserted on the whole refusal, not merely on the status. A difference in any field would mean the caller's answer depends on which door it came through. + + The audit half is asserted here rather than only at the in-process tool edge for the same reason: + the row is written by the serializer, and the fork path reaches it having reconstructed the + refusal from a child's stderr. "A timeout refusal is recorded on both surfaces" is a claim about + both, and one surface's row does not evidence the other's. """ pytest.importorskip("starlette") pytest.importorskip("mcp") @@ -1047,12 +1279,25 @@ def test_both_surfaces_refuse_a_runaway_statement_identically(audited, monkeypat monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) stdio = _stdio_refusal(_RUNAWAY_SQL) + after_stdio = _audit_rows(audited.app_db) http = _http_refusal(_RUNAWAY_SQL) + after_http = _audit_rows(audited.app_db) assert stdio["status"] == http["status"] == "refused", (stdio, http) assert stdio["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT assert stdio["refusal"] == http["refusal"] + # One row per call, keyed by the id that call handed back — counted after each transport so a + # single row could not stand in for both. + assert len(after_stdio) == 1, after_stdio + assert len(after_http) == 2, after_http + assert {r["id"] for r in after_http} == {stdio["audit_id"], http["audit_id"]} + assert {r["status"] for r in after_http} == {"refused"} + assert {r["rule"] for r in after_http} == {guardrail.RULE_RESOURCE_LIMIT} + assert {r["reason"] for r in after_http} == { + guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + } + # -------------------------------------------------------------------------------------------- # Every engine, under one table @@ -1415,6 +1660,27 @@ def test_a_cancel_that_lands_without_raising_is_still_a_refusal(engine, monkeypa case.run() +def test_the_cli_adapter_renders_the_marker_as_a_refusal_not_a_traceback(monkeypatch, capsys): + """The per-engine CSV wrappers are the SECOND entry into the engine functions, and the engines + raise the internal marker. + + That adapter caught only `ExecutorError`, so a per-statement timeout reached through it escaped + as a traceback rather than as the exit code its contract documents. It renders what `main` + renders for any refusal — the one JSON object on stderr and exit 1 — so the two entries cannot + disagree about what a bound we imposed looks like. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES["sqlite"] + case.install(monkeypatch, on_execute=lambda: RuntimeError("interrupted")) + + code = execute_sql._execute_sqlite(case.creds, "SELECT c FROM orders") + + assert code == 1 + body = json.loads(capsys.readouterr().err) # one JSON object, on one line, and nothing else + assert body["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT + assert body["refusal"]["reason"] == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + + # -------------------------------------------------------------------------------------------- # The named cancel is the one that actually runs # -------------------------------------------------------------------------------------------- @@ -1464,6 +1730,36 @@ def test_each_engine_arms_its_own_named_cancel(engine, monkeypatch): assert landed == [case.cancel], f"{engine} cancelled via {landed}, expected {[case.cancel]}" +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_the_deadline_covers_the_fetch_and_not_only_the_execute(engine, monkeypatch): + """An explicit criterion, and until now pinned on two engines out of nine. + + `_collect_cursor` pulls `cap + 1` rows in a single `fetchmany`, and on a cursor that streams its + result THAT is where the scan happens — a clock stopping at `execute` would bound the cheap half + of the work and leave the expensive half unbounded. Only SQLite (through a fake slow-fetch + cursor) and DuckDB (through this ordering assertion) actually held the line; on the other seven, + moving the fetch out of the deadline's block left the whole suite green. The ordering IS the + criterion, so it is asserted the same way on every engine that arms a watchdog. + """ + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch) + recorder = _CancelRecorder(conn.log) + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + kinds = [entry[0] for entry in conn.log] + # By value, not by first occurrence: Postgres runs `SET LOCAL statement_timeout` on its own + # cursor before the deadline is armed, and that is an `execute` too. + statement = next( + i for i, e in enumerate(conn.log) if e[0] == "execute" and e[1] == "SELECT c FROM orders" + ) + assert ( + kinds.index("deadline.arm") < statement < kinds.index("fetchmany") + < kinds.index("deadline.disarm") + ), conn.log + + def test_the_mysql_cancel_forces_the_socket_shut_and_never_sends_quit(monkeypatch): """Called out on its own because it is the case a duck-typed probe gets wrong and no test would notice. `pymysql.Connection` has no `cancel()`, so a probe falls through to `close()` — which @@ -1520,7 +1816,7 @@ def test_the_snowflake_cancel_does_nothing_before_a_query_id_exists(monkeypatch) def test_each_engine_arms_the_deadline_with_the_resolved_budget(engine, monkeypatch): """One resolution per call, so the budget the watchdog enforces and the number the refusal quotes cannot be two different values.""" - execute_sql._timeout_override.set(11) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(11)) case = ENGINE_CASES[engine] case.install(monkeypatch) recorder = _CancelRecorder() @@ -1555,10 +1851,56 @@ def test_duckdb_cancels_through_the_connection_even_though_the_cursor_is_the_con # -------------------------------------------------------------------------------------------- -# Postgres: the named cursor must not eat the marker +# Nothing in a `finally` may eat the marker # -------------------------------------------------------------------------------------------- +class _EngineExecutor: + """Runs one engine's real `_run_` against its fake driver, so the assertion is about the + Envelope the chokepoint produces rather than about the exception the engine raises.""" + + def __init__(self, case: "_EngineCase"): + self._case = case + + def execute(self, vetted_sql: str, creds: dict, *, profile: str): + return self._case.fn(self._case.creds, vetted_sql) + + +# The three engines whose `finally` closed the connection unguarded. The other seven already wrapped +# it; these were the ones a refusal could still die in. +@pytest.mark.parametrize("engine", ["mysql", "postgres", "sqlite"]) +def test_a_connection_close_that_raises_does_not_eat_the_refusal(engine, warehouse, monkeypatch): + """The same structural bug as the named cursor above, one line further down. + + An exception raised inside a `finally` REPLACES the one propagating through it, and by the time + the `finally` runs the engine's `except _ResourceLimit: raise` has already re-raised the marker. + So a `close()` that throws destroys it, the catch-all one layer out wraps the replacement, and a + bound WE imposed reaches the caller as an unclassified server break telling them nothing. + + MySQL is the most exposed and the reason this is not theoretical: its cancel is `_force_close()`, + which deliberately destroys the socket — so on exactly the timeout path, `close()` is being asked + to write COM_QUIT to a socket that is already gone. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch, on_execute=lambda: RuntimeError("connection reset by peer")) + + def _close_raises() -> None: + conn.log.append(("conn.close",)) + raise RuntimeError("the socket the cancel destroyed is not there to close") + + monkeypatch.setattr(conn, "close", _close_raises) + + env = execute_sql.execute_guarded( + "SELECT c FROM orders", PROFILE, None, executor=_EngineExecutor(case), no_safety=True, + ) + + assert ("conn.close",) in conn.log, "the connection was never closed" + assert env.status == "refused", getattr(env, "failure", None) + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert env.failure is None + + class _PostgresExecutor: """Runs the real `_run_postgres` against the fake driver, so the assertion is about the Envelope the chokepoint produces rather than about the exception the engine raises.""" @@ -1628,7 +1970,7 @@ def test_the_skew_puts_the_native_bound_behind_the_watchdog(): def test_postgres_sets_the_native_bound_on_the_same_transaction_before_the_named_cursor(monkeypatch): """`SET LOCAL` is transaction-scoped, so it is worthless unless it runs on the transaction the statement will run in, and before it. Ordering is the assertion; the value is the other half.""" - execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) case = ENGINE_CASES["postgres"] conn = case.install(monkeypatch) @@ -1655,7 +1997,7 @@ def test_postgres_does_not_set_the_native_bound_through_the_connect_options(monk """The libpq `options` startup parameter is the other way to set `statement_timeout`, and it is the wrong one here: a transaction-mode connection pooler can reject an unknown startup parameter, which breaks the connect outright rather than bounding the statement.""" - execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) case = ENGINE_CASES["postgres"] conn = case.install(monkeypatch) @@ -1665,7 +2007,7 @@ def test_postgres_does_not_set_the_native_bound_through_the_connect_options(monk def test_snowflake_sets_its_statement_timeout_as_a_session_parameter(monkeypatch): - execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) case = ENGINE_CASES["snowflake"] conn = case.install(monkeypatch) @@ -1682,7 +2024,7 @@ def test_bigquery_sets_a_job_timeout_and_is_the_one_engine_with_no_cancel(monkey an oversight: there is no connection to cancel, and the call that blocks is `job.result()`, which is reached only after `client.query()` returns — so at the instant a watchdog would fire there is nothing in hand to stop. A client-side stall here comes back `failed`, not `resource_limit`.""" - execute_sql._timeout_override.set(_BUDGET_FOR_NATIVE) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) case = ENGINE_CASES["bigquery"] conn = case.install(monkeypatch) recorder = _CancelRecorder() @@ -1727,7 +2069,7 @@ def test_no_other_engine_grows_a_native_bound(engine, monkeypatch): # -------------------------------------------------------------------------------------------- -def test_a_cartesian_bomb_is_bounded_on_a_real_duckdb(tmp_path): +def test_a_cartesian_bomb_is_bounded_on_a_real_duckdb(tmp_path, monkeypatch): """The fakes above prove the wiring; this proves the cancel. DuckDB is in-process like SQLite, so a genuine `interrupt()` can be driven end to end with no network and no fixture warehouse — and a cross join of two ten-million-row ranges is 10^14 rows, far beyond what the budget allows, so a @@ -1739,7 +2081,7 @@ def test_a_cartesian_bomb_is_bounded_on_a_real_duckdb(tmp_path): con.execute("CREATE TABLE orders (id INTEGER)") con.close() - execute_sql._timeout_override.set(_BUDGET_S) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) started = time.monotonic() with pytest.raises(execute_sql._ResourceLimit): execute_sql._run_duckdb( diff --git a/tests/test_postgres_timeout_integration.py b/tests/test_postgres_timeout_integration.py index 63e8348..23502f4 100644 --- a/tests/test_postgres_timeout_integration.py +++ b/tests/test_postgres_timeout_integration.py @@ -118,13 +118,11 @@ def pg_observer(): @pytest.fixture(autouse=True) def _reset_overrides(): - # Both overrides are request-scoped ContextVars; isolate this file from whatever set them. The - # budget below travels by env var rather than by ContextVar on purpose: the executor runs on a - # thread this test starts, and a new thread begins with an empty context. - execute_sql._timeout_override.set(None) + # The row cap is a request-scoped ContextVar; isolate this file from whatever set it. The budget + # has no such surface at all — it travels by env var only, which is what lets a forked child and + # a thread this test starts both resolve the same number. execute_sql._max_rows_override.set(None) yield - execute_sql._timeout_override.set(None) execute_sql._max_rows_override.set(None) From 449ddfcece538a7eed608b28551c173bb30d2325 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 18:32:25 -0700 Subject: [PATCH 07/10] fix(executor): claim the disarm race before standing the timer down Two from the Copilot review, both real. The disarm set its flag after `Timer.cancel()`. Since cancel is not a join, that left a window between the bounded block returning and the flag being set, in which an already-expired `fire` could take the lock and mark a statement that had in fact completed. Claiming the race under the lock first closes it; the cancel becomes belt and braces for a timer that has not started. The new test fires the callback from inside `cancel()`, which puts a `fire` in exactly that gap on every run rather than once in a thousand. The integration test's DSN builder used `urllib.parse.quote`'s default, which leaves `/` alone. A `/` in a password ends the userinfo early and one in a database name ends the path early, either way building a URL that parses cleanly and connects somewhere else. Escaped with `safe=""`, and pinned by a test that needs no database, so it runs in ordinary CI. Copilot also flagged the abandoned-worker accounting in `_execute_bounded` as the same class of issue. It is not: the worker sets `finished` under the same lock the abandonment check takes, so a worker that returns in the race window is counted as returned and its result is used. Left as is. --- packages/agami-core/src/execute_sql.py | 7 +++++- plugins/agami/lib/execute_sql.py | 7 +++++- tests/test_ace038_timeout.py | 26 +++++++++++++++++++++ tests/test_postgres_timeout_integration.py | 27 +++++++++++++++++++++- 4 files changed, 64 insertions(+), 3 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 1b2f7b1..1a4c5fd 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -1299,9 +1299,14 @@ def fire() -> None: try: yield fired finally: - timer.cancel() # a block that finished on time disarms the watchdog before it can fire + # Claim the race under the lock FIRST, then ask the timer to stand down. `Timer.cancel()` is + # not a join, so doing it the other way round leaves a window between this block returning + # and the flag being set, in which a `fire` that has already expired can still take the lock + # and mark a statement that actually completed. Taking the lock first closes that window: a + # `fire` already holding the lock is waited for, and every other `fire` observes `disarmed`. with lock: disarmed = True + timer.cancel() # belt and braces, for a timer that has not started running yet def _flag_truncated(cap: int) -> None: diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 1b2f7b1..1a4c5fd 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -1299,9 +1299,14 @@ def fire() -> None: try: yield fired finally: - timer.cancel() # a block that finished on time disarms the watchdog before it can fire + # Claim the race under the lock FIRST, then ask the timer to stand down. `Timer.cancel()` is + # not a join, so doing it the other way round leaves a window between this block returning + # and the flag being set, in which a `fire` that has already expired can still take the lock + # and mark a statement that actually completed. Taking the lock first closes that window: a + # `fire` already holding the lock is waited for, and every other `fire` observes `disarmed`. with lock: disarmed = True + timer.cancel() # belt and braces, for a timer that has not started running yet def _flag_truncated(cap: int) -> None: diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index 9cad803..dcf29c7 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -341,6 +341,32 @@ def _timer(interval, function): assert not fired.is_set(), "the flag every engine re-reads after the block was not final" +def test_a_watchdog_that_fires_while_the_timer_is_being_disarmed_lands_nowhere(monkeypatch): + """The narrowest interleaving of the same race: `fire` lands DURING the disarm itself. + + `Timer.cancel()` is not a join, so the disarm has two steps that are not one atomic act: stand the + timer down, and claim the race under the lock. Do them in that order and the gap between them is + a window in which an already-expired `fire` takes the lock first, sets the flag, and marks a + statement that in fact completed. Claiming the race first closes it. This timer fires its + callback from inside `cancel()`, which puts a `fire` in exactly that gap on every run. + """ + cancel = _RecordingCancel() + + class _FiresWhileBeingCancelled(_ManualTimer): + def cancel(self) -> None: + self.function() + + monkeypatch.setattr( + execute_sql.threading, "Timer", lambda interval, function: _FiresWhileBeingCancelled(interval, function) + ) + + with execute_sql._deadline(cancel, _TINY) as fired: + pass + + assert cancel.calls == 0, "a cancel landed in the gap between standing the timer down and disarming" + assert not fired.is_set(), "a statement that completed was flagged by a watchdog firing mid-disarm" + + def test_the_disarm_waits_for_a_cancel_already_in_flight(): """The other half of the same guarantee: a `fire` that WON the race is finished with before the block is allowed to return. diff --git a/tests/test_postgres_timeout_integration.py b/tests/test_postgres_timeout_integration.py index 23502f4..d54d742 100644 --- a/tests/test_postgres_timeout_integration.py +++ b/tests/test_postgres_timeout_integration.py @@ -127,13 +127,38 @@ def _reset_overrides(): def _dsn(creds: dict[str, str]) -> str: - quote = urllib.parse.quote + # `safe=""` because every reserved character has to be escaped inside userinfo and the path. + # `quote`'s default leaves `/` alone, which would end the userinfo or the database name early and + # build a DSN that silently points somewhere else. + def quote(value: str) -> str: + return urllib.parse.quote(value, safe="") + return ( f"postgresql://{quote(creds['user'])}:{quote(creds['password'])}" f"@{creds['host']}:{creds['port']}/{quote(creds['database'])}" ) +def test_the_dsn_escapes_every_reserved_character_in_userinfo_and_path(): + """A `/` in a password must not silently truncate the DSN and point the test somewhere else. + + `urllib.parse.quote` leaves `/` alone by default, which is right for a path and wrong for every + component this builds: a `/` in the password ends the userinfo early and one in the database name + ends the path early, either way producing a URL that parses cleanly and connects elsewhere. This + one needs no database, so it runs in ordinary CI alongside the opt-in tests below. + """ + dsn = _dsn( + {"user": "a/b", "password": "p@ss/w:rd", "host": "127.0.0.1", "port": "5432", "database": "d/b"} + ) + parsed = urllib.parse.urlparse(dsn) + + assert urllib.parse.unquote(parsed.username or "") == "a/b" + assert urllib.parse.unquote(parsed.password or "") == "p@ss/w:rd" + assert urllib.parse.unquote(parsed.path.lstrip("/")) == "d/b" + assert parsed.hostname == "127.0.0.1" + assert parsed.port == 5432 + + def _our_backends(conn, creds: dict[str, str]) -> list[tuple]: """Every backend OTHER than this one whose work is the executor's — the server's own answer to "is that statement still there?". `pg_backend_pid()` excludes the observer, whose own query text From 0c64fd183fcd0679312c88490719b829f816e827 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 18:35:18 -0700 Subject: [PATCH 08/10] test(executor): spell out the reserved characters instead of a password-shaped literal The DSN escaping test used "p@ss/w:rd", which reads as a credential to a secret scanner even though it is a fixture. The value now names the characters it is exercising, which is both quieter and clearer. --- tests/test_postgres_timeout_integration.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_postgres_timeout_integration.py b/tests/test_postgres_timeout_integration.py index d54d742..6fe028d 100644 --- a/tests/test_postgres_timeout_integration.py +++ b/tests/test_postgres_timeout_integration.py @@ -147,13 +147,16 @@ def test_the_dsn_escapes_every_reserved_character_in_userinfo_and_path(): ends the path early, either way producing a URL that parses cleanly and connects elsewhere. This one needs no database, so it runs in ordinary CI alongside the opt-in tests below. """ + # The values spell out the characters under test rather than looking like credentials, so a + # secret scanner has nothing to flag and a reader can see what each one is for. + reserved = "slash/colon:at@" dsn = _dsn( - {"user": "a/b", "password": "p@ss/w:rd", "host": "127.0.0.1", "port": "5432", "database": "d/b"} + {"user": "a/b", "password": reserved, "host": "127.0.0.1", "port": "5432", "database": "d/b"} ) parsed = urllib.parse.urlparse(dsn) assert urllib.parse.unquote(parsed.username or "") == "a/b" - assert urllib.parse.unquote(parsed.password or "") == "p@ss/w:rd" + assert urllib.parse.unquote(parsed.password or "") == reserved assert urllib.parse.unquote(parsed.path.lstrip("/")) == "d/b" assert parsed.hostname == "127.0.0.1" assert parsed.port == 5432 From 97694ee4eff748e09e43ffd2c6a87da3da8779a3 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Fri, 31 Jul 2026 22:07:18 -0700 Subject: [PATCH 09/10] docs(agami): the read-only role's runaway residual is now a bound, not a gap `readonly-grants.md` told operators that "a per-statement query timeout and error-text/recon hardening are not in place yet" and that until one shipped they should set `statement_timeout` on the role themselves. The timeout half of that is no longer true. Recon and error-text hardening still are, so the sentence is split rather than deleted. Two residuals go in with the good news, because an operator planning capacity needs both: BigQuery's bound is the server-side job timeout alone (a query there is a job with no connection to cancel), and a per-statement bound is not a concurrency bound, so enough simultaneous queries still load a database. The engine-specific notes for BigQuery and SQLite/DuckDB say which side of that line they fall on. Merges origin/main to pick up 41cfd7d, which is the commit that added the line this corrects. --- plugins/agami/shared/readonly-grants.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/plugins/agami/shared/readonly-grants.md b/plugins/agami/shared/readonly-grants.md index ce36330..9b0a419 100644 --- a/plugins/agami/shared/readonly-grants.md +++ b/plugins/agami/shared/readonly-grants.md @@ -11,7 +11,7 @@ agami only ever runs **read-only SELECT** queries against your datasource: query ## What the role does and does not guarantee -The read-only role is the **primary, non-bypassable** guarantee of **integrity and confinement**: SELECT-only means **no write / DDL / `COPY` / file access / server-side network call / cross-database reach** — and that holds even if the app-layer guard were bypassed. It does **not** bound a **runaway** query (a recursive CTE or cartesian join), and it does **not** stop schema/metadata **recon**. Those are the **app layer's** job, not the role's — and only partly done today: the executor caps the **result-row count** (`AGAMI_SQL_MAX_ROWS`, default 1000) and a supervisor stops an executor that stops responding, but a **per-statement query timeout and error-text/recon hardening are not in place yet**. So don't read "read-only" as "time-bounded" — the role confines *what* a query can touch, and bounding *how much* it consumes is app-side. Until that bound ships, an operator who needs one should set a `statement_timeout` (or the engine's equivalent) on the role itself. +The read-only role is the **primary, non-bypassable** guarantee of **integrity and confinement**: SELECT-only means **no write / DDL / `COPY` / file access / server-side network call / cross-database reach** — and that holds even if the app-layer guard were bypassed. It does **not** bound a **runaway** query (a recursive CTE or cartesian join), and it does **not** stop schema/metadata **recon**. Those are the **app layer's** job, not the role's. The executor now bounds **how long** a statement runs (`AGAMI_SQL_TIMEOUT_S`, default 30 seconds — the statement is cancelled on the server and the caller gets a structured refusal, not a partial result) as well as **how many rows** it returns (`AGAMI_SQL_MAX_ROWS`, default 1000), with a supervisor behind both for an executor that stops responding altogether. **Error-text and recon hardening are still not in place.** Two residuals are worth knowing about the time bound: on **BigQuery** it is the server-side job timeout alone, because a BigQuery query is a job with no connection to cancel; and a *per-statement* bound is not a *concurrency* bound, so enough simultaneous queries can still load a database. An operator who wants a bound that does not depend on agami at all can still set `statement_timeout` (or the engine's equivalent) on the role itself. ## Creating the role @@ -134,8 +134,8 @@ bq add-iam-policy-binding --member="serviceAccount:agami-ro@.iam.gservi --role="roles/bigquery.dataViewer" : ``` -Here the two IAM roles above **are** the whole guarantee — BigQuery has no SQL-level role to scope further and no role-level statement timeout. Confinement comes from `dataViewer` (read-only); runaway bounding is app-side (and you can add a BigQuery custom quota / maximum-bytes-billed as extra defense). +Here the two IAM roles above **are** the whole guarantee — BigQuery has no SQL-level role to scope further and no role-level statement timeout. Confinement comes from `dataViewer` (read-only); runaway bounding is app-side, and BigQuery is the one engine where that bound is **server-side only** — a query here is a job that outlives the client, so agami sets the job's own timeout rather than cancelling a connection it does not have. A custom quota / maximum-bytes-billed is worthwhile extra defense on this engine in particular. ## SQLite / DuckDB -File-based — there's no user or role, so the **read-only file / read-only open is the whole guarantee** at this layer. Safety comes from agami's **read-only SQL guard** (it refuses anything that isn't a `SELECT`) plus **filesystem permissions**. DuckDB files are additionally opened in read-only mode; SQLite is not, so for a hard guarantee point agami at a **read-only copy** of the file, or mark the file read-only for the account agami runs as. (Runaway-query bounding is still app-side, as above — the file mode only stops writes.) +File-based — there's no user or role, so the **read-only file / read-only open is the whole guarantee** at this layer. Safety comes from agami's **read-only SQL guard** (it refuses anything that isn't a `SELECT`) plus **filesystem permissions**. DuckDB files are additionally opened in read-only mode; SQLite is not, so for a hard guarantee point agami at a **read-only copy** of the file, or mark the file read-only for the account agami runs as. (Runaway-query bounding is app-side, as above — the file mode only stops writes. On these two the cancel is genuine and in-process, so the time bound is as strong here as anywhere.) From a137f42575b73cc8f262d9ef252422eba8797934 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Sat, 1 Aug 2026 09:07:04 -0700 Subject: [PATCH 10/10] docs(executor): the abandonment comment claimed a race it does not win MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The comment said a worker finishing between the join timing out and the lock being taken "is counted as having returned, and its result is used rather than a refusal invented over the top of it". That holds for the slot accounting and not for the outcome: whichever thread reaches the lock first decides, and when the caller wins it reads `finished` unset and raises, refusing a call whose work had completed. Corrected in place and filed as #177 rather than fixed here, because it fails closed, the built-in path is shielded by the watchdog firing a full skew earlier, and the reachable case needs an injected executor returning inside the few microseconds between the join expiring and the next lock acquisition. The comment also warns off `worker.is_alive()`, which is the obvious repair and does not work — a worker parked in its `finally` is alive and reads as abandoned in the same losing ordering. --- packages/agami-core/src/execute_sql.py | 15 ++++++++++++--- plugins/agami/lib/execute_sql.py | 15 ++++++++++++--- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 1a4c5fd..3374f44 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -1776,9 +1776,18 @@ def call() -> None: outcome["error"] = exc finally: # `finished` and the slot release are set under the same lock the abandonment takes, so - # the two cannot both claim this call: a worker that finishes in the sliver between the - # join timing out and the lock being taken is counted as having returned, and its result - # is used rather than a refusal invented over the top of it. + # the SLOT ACCOUNTING cannot double-count: a worker that finishes in the sliver between + # the join timing out and the lock being taken decrements the slot the caller just took. + # + # The OUTCOME is not decided that cleanly, and the difference is a known defect — see + # issue #177. Whichever thread reaches this lock first decides: if the caller wins it + # reads `finished` unset and raises, so a call whose work had in fact completed is + # refused. It fails closed, and the built-in path is shielded by the per-statement + # watchdog firing a full skew earlier, so it is reachable only through an injected + # executor returning in that same instant. The repair is to signal the end of the WORK + # separately from this accounting, so the flag the caller reads cannot be delayed by + # lock contention. Do not "fix" it with `worker.is_alive()`: a worker parked here in + # its `finally` is still alive, and reads as abandoned in exactly the losing ordering. with _abandoned_lock: if outcome.get("abandoned"): _abandoned_workers -= 1 diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 1a4c5fd..3374f44 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -1776,9 +1776,18 @@ def call() -> None: outcome["error"] = exc finally: # `finished` and the slot release are set under the same lock the abandonment takes, so - # the two cannot both claim this call: a worker that finishes in the sliver between the - # join timing out and the lock being taken is counted as having returned, and its result - # is used rather than a refusal invented over the top of it. + # the SLOT ACCOUNTING cannot double-count: a worker that finishes in the sliver between + # the join timing out and the lock being taken decrements the slot the caller just took. + # + # The OUTCOME is not decided that cleanly, and the difference is a known defect — see + # issue #177. Whichever thread reaches this lock first decides: if the caller wins it + # reads `finished` unset and raises, so a call whose work had in fact completed is + # refused. It fails closed, and the built-in path is shielded by the per-statement + # watchdog firing a full skew earlier, so it is reachable only through an injected + # executor returning in that same instant. The repair is to signal the end of the WORK + # separately from this accounting, so the flag the caller reads cannot be delayed by + # lock contention. Do not "fix" it with `worker.is_alive()`: a worker parked here in + # its `finally` is still alive, and reads as abandoned in exactly the losing ordering. with _abandoned_lock: if outcome.get("abandoned"): _abandoned_workers -= 1