From ca57e2cc1fb234fda8df3ae30e7f9378046e95f1 Mon Sep 17 00:00:00 2001 From: Ashwin Ramachandran Date: Mon, 14 Sep 2026 15:32:11 -0700 Subject: [PATCH 1/3] Let an organisation have its own row cap and statement time limit (engine seam) Core stores no setting. An embedder registers a provider (org_id) -> {"max_rows", "timeout_s"} via Adapters.statement_limits or tools.set_statement_limits_provider; missing or unusable values fall back to AGAMI_SQL_MAX_ROWS / AGAMI_SQL_TIMEOUT_S. No ceiling. The limits are resolved once per execute_sql call into a ContextVar both resolvers read before the environment, and _pass_child_env writes the same numbers into the forked child's AGAMI_SQL_* keys, so every bound in the ordered family derives from one budget on both sides of the fork. statement_limits() reports the caller's organisation's limits, statement_limit_defaults() the deployment and recommended values, and the execute_sql description states the caller's numbers at list-tools time. Refs #329 Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 19 ++ packages/agami-core/src/execute_sql.py | 117 ++++++-- packages/agami-core/src/mcp_harness.py | 8 +- packages/agami-core/src/mcp_http.py | 25 +- packages/agami-core/src/ports.py | 12 +- packages/agami-core/src/tools.py | 231 +++++++++++++-- plugins/agami/lib/execute_sql.py | 117 ++++++-- tests/test_ace038_timeout.py | 17 +- tests/test_per_org_statement_limits.py | 391 +++++++++++++++++++++++++ 9 files changed, 843 insertions(+), 94 deletions(-) create mode 100644 tests/test_per_org_statement_limits.py diff --git a/CHANGELOG.md b/CHANGELOG.md index deb169b1..2bcfa463 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,25 @@ below corresponds to one such version. ## [Unreleased] +### Added + +- **An organisation can have its own row cap and statement time limit** (#329, engine half). Core + stores no such setting: an embedder registers a provider, `(org_id) -> {"max_rows", "timeout_s"}`, + through `Adapters.statement_limits` or `tools.set_statement_limits_provider`. A missing, `None` or + unusable value (not a positive whole number, or a provider that raises) falls back to + `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`, which stay the deployment default, with a warning in + the log. There is no ceiling. + - The limits are resolved once per `execute_sql` call and held for the whole call, and the forked + child is handed the same two numbers in its environment, so the watchdog, the native bound, the + outer bound and the supervisor still derive from one budget on both sides of the fork. + - `tools.statement_limits(org_id=None)` reports the limits in force for the current (or a named) + organisation; `tools.statement_limit_defaults()` reports the deployment values and the + recommended ones (1000 rows, 30 seconds). `tools.pinned_statement_limits(org_id)` lets a direct + caller of `execute_guarded` apply an organisation's limits. + - The `execute_sql` description states the caller's organisation's numbers, built when tools are + listed rather than once at start-up. A client keeps the list for its session, so a changed limit + reaches new sessions; an existing session meets it in the refusal, which names the number per call. + ### Fixed - **A table outside the connection's default schema resolves when the client names it without its diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 896bb7c6..8b3f5126 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -74,7 +74,7 @@ import urllib.parse import uuid from collections.abc import Callable, Iterator -from contextvars import ContextVar, copy_context +from contextvars import ContextVar, Token, copy_context from dataclasses import asdict, dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -1514,19 +1514,39 @@ def _pin_model_pass_posture() -> bool: return value -def _resolve_row_cap() -> int: - """Effective result-row cap. `AGAMI_SQL_MAX_ROWS` is the operator-configurable DEPLOYMENT cap - (default 1000 when unset) — an operator owns their availability tradeoff and may set it higher OR - lower than 1000; it is NOT a hard 1000 ceiling. A missing/invalid/zero env value falls back to - 1000. - - The operator is the only voice here. A per-call override used to be able to lower it, and it went - with the trim (ACE-087): the one thing a caller might know better than the deployment — that it - wants MORE rows — is the thing a lowering-only override structurally could not express, and a - caller that wants 200 rows says so in the statement, where the intent is legible to everything - downstream.""" +# The effective (row cap, timeout seconds) for THIS call, when the caller's organisation has its own +# (#329). Absent means "the deployment's environment", which is every call that did not come through +# a pinning entry point. Both numbers in one value rather than two ContextVars, so a reader can never +# see one organisation's row cap beside another's deadline. +# +# This is a second, higher-precedence input to the budget, which is exactly the hazard +# `_resolve_timeout_s` used to rule out by having no such thing: a value that outranks the environment +# in the parent and is invisible to a forked child would let the supervisor bound the parent derives +# sit below the budget the child enforces. It is allowed now on the same terms `_pass_posture` was: +# it is resolved ONCE per call, before any bound is derived, and `tools._pass_child_env` writes both +# numbers into the child's `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`, so the child re-resolves the +# identical budget from its environment. The fork carries it explicitly; it is never lost across it. +# +# The values held here are already validated positive ints — the provider that supplies them lives in +# `tools`, which does the checking, so nothing here needs to re-parse them. +_statement_limits: ContextVar[tuple[int, int] | None] = ContextVar( + "_statement_limits", default=None +) + + +def _pin_statement_limits(max_rows: int, timeout_s: int) -> Token[tuple[int, int] | None]: + """Fix this call's effective budget. Returns the token the caller must reset with, in a + `finally`: the pin is request-scoped, and one left behind on a reused worker would hand the next + organisation's call this one's limits.""" + return _statement_limits.set((max_rows, timeout_s)) + + +def _row_cap_from_env() -> int: + """The DEPLOYMENT row cap: `AGAMI_SQL_MAX_ROWS`, default 1000 when unset. An operator owns their + availability tradeoff and may set it higher OR lower than 1000; it is NOT a hard 1000 ceiling. A + missing/invalid/zero env value falls back to 1000.""" raw = os.environ.get("AGAMI_SQL_MAX_ROWS", "").strip() - # `isdecimal`, not `isdigit`, for the reason `_resolve_timeout_s` gives: `isdigit` admits `²`, + # `isdecimal`, not `isdigit`, for the reason `_timeout_s_from_env` gives: `isdigit` admits `²`, # which `int()` then refuses. `tools` now resolves this while building its registry, so a raise # here would stop the module importing at all rather than failing one call. cap = int(raw) if raw.isdecimal() else _DEFAULT_MAX_ROWS @@ -1535,6 +1555,22 @@ def _resolve_row_cap() -> int: return cap +def _resolve_row_cap() -> int: + """Effective result-row cap for THIS call: the organisation's own when a call pinned one + (`_statement_limits`, #329), otherwise the deployment's (`_row_cap_from_env`). + + No per-CALL override exists, and that is unchanged. One used to be able to lower the cap, and it + went with the trim (ACE-087): the one thing a caller might know better than the deployment — that + it wants MORE rows — is the thing a lowering-only override structurally could not express, and a + caller that wants 200 rows says so in the statement, where the intent is legible to everything + downstream. What the pin carries is an administrator's setting for the organisation, resolved + before the call and not chosen by it.""" + pinned = _statement_limits.get() + if pinned is not None: + return pinned[0] + return _row_cap_from_env() + + _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 @@ -1581,18 +1617,34 @@ def _resolve_row_cap() -> int: 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. - - **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 + """Effective per-statement timeout for THIS call, in whole seconds: the organisation's own when + a call pinned one (`_statement_limits`, #329), otherwise the deployment's (`_timeout_s_from_env`). + Every bound in the ordered family — watchdog, native skew, outer bound, supervisor — reads this, + so they all derive from the one effective budget. + + **Two sources, and one budget on both sides of the fork.** The environment used to be 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. A per-organisation limit needs exactly such an override, + so the hazard is now closed by construction instead of by absence: the pin is set once per call, + before the supervisor bound is derived, and `tools._pass_child_env` writes the pinned numbers into + the child's environment, where this same resolver (with no pin of its own) reads them back. The + parent and the child therefore reach the identical number, which is the property the old rule + existed to guarantee.""" + pinned = _statement_limits.get() + if pinned is not None: + return pinned[1] + return _timeout_s_from_env() + + +def _timeout_s_from_env() -> int: + """The DEPLOYMENT per-statement timeout: `AGAMI_SQL_TIMEOUT_S`, 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. It is also what an organisation with no limit of its own gets, + and what a forked child reads the pinned budget back from. + + Unlike `_row_cap_from_env`, 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 @@ -1711,8 +1763,9 @@ def _resource_limit_refusal(exc: _ResourceLimit | None) -> Refusal: invariant that survives is one rule with one emit site, not one sentence. The budget is re-resolved rather than carried: nothing between the engine call and here can - change the environment the resolvers read, so both `_resolve_timeout_s` and `_resolve_row_cap` - return the same number the bound itself used. The configured number belongs in the detail — it + change what the resolvers read — the call's pinned organisation limits, or the environment when + none were pinned (and, in a forked child, the environment its parent wrote the pin into) — so + both `_resolve_timeout_s` and `_resolve_row_cap` return the same number the bound itself 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. """ @@ -2636,8 +2689,9 @@ def _execute_bounded( The call runs inside a copy of the CALLER's context, and what that is FOR changed when the per-call row cap went (ACE-087). It used to carry ``_max_rows_override`` to ``_resolve_row_cap`` - inside the worker; the cap is the deployment's environment now and needs no carrier. What it - still carries is the *caller's* request scope — ``tools._current_org_ctx``, the resolve-once + inside the worker; that override is gone. It now also carries the call's pinned organisation + limits (``_statement_limits``, #329), though nothing in the worker reads them today — every bound + is derived on the caller's side. What it chiefly carries is the *caller's* request scope — ``tools._current_org_ctx``, the resolve-once request cache, the actor and session on the served path — into the one place a consumer's own code runs. That is the point of the ``Executor`` seam: a pooled / per-user-RBAC executor picks its connection from exactly that context, and a new thread starts with an empty one, so dropping @@ -2742,8 +2796,9 @@ def execute_guarded( ``_load_credentials`` sits INSIDE the try deliberately, so a bad profile / missing DSN becomes a ``failed``/``dsn`` Envelope carrying its detailed message rather than escaping as an exception - the two callers would each have to translate. The row cap is the deployment's alone - (``AGAMI_SQL_MAX_ROWS``); no caller can lower it for one call.""" + the two callers would each have to translate. The row cap is the organisation's + limit when the caller pinned one (``tools.pinned_statement_limits``) and the deployment's + ``AGAMI_SQL_MAX_ROWS`` otherwise; no caller can change it for one statement.""" # Clear before anything can set it, so a detail from a PREVIOUS call in this context can never # be attributed to this one. The recorder reads it unconditionally; a stale value would put the # wrong error text on a row that succeeded. diff --git a/packages/agami-core/src/mcp_harness.py b/packages/agami-core/src/mcp_harness.py index fdbd9482..57dc305e 100644 --- a/packages/agami-core/src/mcp_harness.py +++ b/packages/agami-core/src/mcp_harness.py @@ -39,6 +39,7 @@ bootstrap_paths, server_instructions, server_version, + tool_description, ) # MCP negotiates a protocol version during `initialize`: the client names the version it wants and @@ -80,7 +81,12 @@ def _handle_initialize(req_id: Any, params: dict[str, Any]) -> None: def _handle_tools_list(req_id: Any) -> None: _result(req_id, { "tools": [ - {"name": name, "description": meta["description"], "inputSchema": meta["inputSchema"]} + # Per caller, as the HTTP server does it (#329), so the two servers state the same numbers. + { + "name": name, + "description": tool_description(name, meta["description"]), + "inputSchema": meta["inputSchema"], + } for name, meta in TOOLS.items() ] }) diff --git a/packages/agami-core/src/mcp_http.py b/packages/agami-core/src/mcp_http.py index 51e4f9da..a5d83b1d 100644 --- a/packages/agami-core/src/mcp_http.py +++ b/packages/agami-core/src/mcp_http.py @@ -58,7 +58,9 @@ server_instructions, server_version, set_injected_executor, + set_statement_limits_provider, thread_id_is_required, + tool_description, typed_outcome_overrides, ) @@ -482,14 +484,27 @@ def _visible(name: str) -> bool: instructions = f"{instructions}\n{extra_instructions}" server = Server(SERVER_NAME, version=server_version(), instructions=instructions) - @server.list_tools() - async def _list_tools() -> list: + def _listed() -> list: return [ - mt.Tool(name=name, description=meta["description"], inputSchema=meta["inputSchema"]) + # `tool_description` states execute_sql's limits for THIS caller's organisation (#329). + # That is core describing its own tool per request, not the subtractive-only hook above + # reshaping one: a consumer's description passes through it untouched. + mt.Tool( + name=name, + description=tool_description(name, meta["description"]), + inputSchema=meta["inputSchema"], + ) for name, meta in registry.items() if _visible(name) ] + @server.list_tools() + async def _list_tools() -> list: + # Off the loop, for ACE-048's reason: the limits provider is the consumer's code and usually a + # database read, and on the loop one slow read would stall every in-flight request. + # `run_blocking` copies the request context, so the organisation is still set in the worker. + return await run_blocking(_listed) + @server.call_tool() async def _call_tool(name: str, arguments: dict) -> list: meta = registry.get(name) @@ -635,6 +650,10 @@ def create_app( # AH-012: register the composition-root executor (None = the default subprocess path). Behind the # shared guard in `tool_execute_sql`; a hosted consumer injects a pooled/RBAC/tunnel executor here. set_injected_executor(adapters.executor) + # #329: register the per-organisation statement-limits provider the same way, and unconditionally + # for the same reason — the adapters are the composition root, so an app built without one must not + # inherit a provider an earlier app in the same process installed. + set_statement_limits_provider(getattr(adapters, "statement_limits", None)) # Validate consumer-supplied tools up front so a malformed entry fails at construction with a # clear error, not later as a KeyError/500 inside tools/list or tools/call. for tool_name, meta in (extra_tools or {}).items(): diff --git a/packages/agami-core/src/ports.py b/packages/agami-core/src/ports.py index a72b1354..6bf55aa6 100644 --- a/packages/agami-core/src/ports.py +++ b/packages/agami-core/src/ports.py @@ -22,9 +22,9 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: # Only for type-checkers — kept out of the runtime import graph so the Protocols (and a @@ -166,3 +166,11 @@ class Adapters: # only — it filters the one shared registry and never adds or reshapes a tool. See # `mcp_http.build_server` for where it is applied (both list AND call, deliberately). tool_visibility: Callable[[str], bool] | None = None + # `statement_limits(org_id) -> {"max_rows": int | None, "timeout_s": int | None} | None` supplies an + # organisation's own row cap and statement deadline (#329). None (the default) is today's + # behaviour: every organisation gets `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`. Core stores no + # such setting — the consumer owns the storage and the admin screen, core only asks — and a + # missing, None or unusable value falls back to the environment. See + # `tools.set_statement_limits_provider` for the contract and `tools.pinned_statement_limits` for how + # one call's answer is held identical on both sides of the fork. + statement_limits: Callable[[str], Mapping[str, Any] | None] | None = None diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index 392a8008..c781ba15 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -20,6 +20,7 @@ from __future__ import annotations +import contextlib import csv import functools import io @@ -32,7 +33,7 @@ import threading import time import uuid -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping from contextvars import ContextVar, Token from dataclasses import asdict from pathlib import Path @@ -1889,6 +1890,116 @@ def set_injected_executor(executor: Any | None) -> None: _INJECTED_EXECUTOR = executor +# The composition-root statement-limits provider (#329). ``None`` (the default) means every +# organisation gets the deployment's own `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`. Core stores no +# per-organisation setting; the consumer that owns the admin screen and its storage registers a +# callable here, and core only asks it. Process-global for the reason `_INJECTED_EXECUTOR` is: the +# provider is a composition-root singleton, while the ANSWER it gives is per-request. +_STATEMENT_LIMITS_PROVIDER: Callable[[str], Mapping[str, Any] | None] | None = None + +# The two keys a provider may answer, and the only two read from its mapping — anything else it +# returns is ignored rather than rejected, so a consumer can carry its own bookkeeping in the row. +_STATEMENT_LIMIT_KEYS = ("max_rows", "timeout_s") + + +def set_statement_limits_provider( + provider: Callable[[str], Mapping[str, Any] | None] | None, +) -> None: + """Register (or clear) the per-organisation statement-limits provider. + + ``provider(org_id)`` returns ``{"max_rows": int | None, "timeout_s": int | None}`` or ``None``. A + missing key, a ``None``, or a ``None`` result means "this organisation has no limit of its own" and + the deployment's environment value applies. There is no ceiling: an administrator may set either + number to any positive whole number. Called by ``mcp_http.create_app`` from + ``adapters.statement_limits``; an embedder with no HTTP app may call it directly.""" + global _STATEMENT_LIMITS_PROVIDER + if provider is not None and not callable(provider): + # Fail at registration, like `set_injected_executor`: a malformed adapter should stop the app + # being built, not surface as a warning on every statement. + raise TypeError("statement limits provider must be callable: (org_id) -> mapping | None") + _STATEMENT_LIMITS_PROVIDER = provider + + +def _provider_limit(org_id: str, key: str, value: Any) -> int | None: + """One provider value, validated: a positive int, or ``None`` to fall back to the deployment. + + ``bool`` is refused although it is an ``int``, because ``True`` reaching the row cap as ``1`` is a + storage bug that would look like a setting. Anything unusable is logged at warning and declined — + never raised: this runs at the entry of every statement and while tools are listed, and one + organisation's bad row must cost that organisation its override, not everybody their query.""" + if value is None: + return None + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value + _LOG.warning( + "statement limits provider returned %s=%r for org %s, which is not a positive whole number; " + "using the deployment value.", + key, + value, + org_id, + ) + return None + + +def _effective_statement_limits(org_id: str | None) -> tuple[int, int]: + """(row cap, timeout seconds) for ``org_id`` (the current request's when None): the provider's + value where it gave a usable one, the deployment's environment value everywhere else. + + The organisation is looked up only when a provider is registered. Without one the answer cannot + depend on it, and the registry renders the limits sentence at import, where resolving an org would + read configuration files for a number that is not going to be used.""" + from execute_sql import _row_cap_from_env, _timeout_s_from_env + + limits = {"max_rows": _row_cap_from_env(), "timeout_s": _timeout_s_from_env()} + provider = _STATEMENT_LIMITS_PROVIDER + if provider is None: + return limits["max_rows"], limits["timeout_s"] + org_id = org_id or _current_org_id() + try: + supplied = provider(org_id) + except Exception: + # The provider is the consumer's code, usually a database read. Failing it must not fail the + # statement: the deployment's own limits are a safe, known answer, and the log carries why. + _LOG.warning( + "statement limits provider failed for org %s; using the deployment values.", + org_id, + exc_info=True, + ) + supplied = None + if supplied is not None and not isinstance(supplied, Mapping): + _LOG.warning( + "statement limits provider returned %r for org %s, not a mapping; using the deployment " + "values.", + type(supplied).__name__, + org_id, + ) + supplied = None + for key in _STATEMENT_LIMIT_KEYS: + value = _provider_limit(org_id, key, (supplied or {}).get(key)) + if value is not None: + limits[key] = value + return limits["max_rows"], limits["timeout_s"] + + +@contextlib.contextmanager +def pinned_statement_limits(org_id: str | None = None) -> Iterator[tuple[int, int]]: + """Resolve an organisation's effective limits ONCE and hold them for the enclosed block. + + ``tool_execute_sql`` opens this around every call; an embedder calling + ``execute_sql.execute_guarded`` directly (an evaluation run, say) opens it itself to get the same + organisation's limits rather than the deployment's. ``org_id`` defaults to the current request's. + Reset in a ``finally``, because a worker thread serves one organisation after another and a pin + that outlived its call would hand the next one these limits.""" + import execute_sql + + max_rows, timeout_s = _effective_statement_limits(org_id) + token = execute_sql._pin_statement_limits(max_rows, timeout_s) + try: + yield max_rows, timeout_s + finally: + execute_sql._statement_limits.reset(token) + + def _finalize_execution( columns: list, data_rows: list, @@ -2350,18 +2461,30 @@ def tool_execute_sql(args: dict[str, Any]) -> str: """ cache_token = begin_request_cache() try: - return _tool_execute_sql(args) + # The organisation's limits are pinned here for the same "one point both paths pass through" + # reason (#329): the in-process path's bounds, the fork path's supervisor bound and the child + # environment `_pass_child_env` builds all read this one resolution, so a provider whose answer + # changes mid-call cannot give the two sides of the fork different budgets. + with pinned_statement_limits(): + return _tool_execute_sql(args) finally: end_request_cache(cache_token) def _pass_child_env() -> dict[str, str]: - """The child's environment: this process's, with the ACE-101 posture written in explicitly. + """The child's environment: this process's, with the ACE-101 posture and this call's effective + statement limits written in explicitly. Everything else is inherited untouched, which the fork depends on: the child re-resolves its own timeout, row cap and credentials from the environment, and the supervisor bound computed on this - side is only correct because the child reaches the identical number. This adds one key and - overrides nothing else. + side is only correct because the child reaches the identical number. + + The two limit keys are why that still holds with a per-organisation limit (#329). The child has no + provider and no pin; it reads `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S` and nothing else. Writing + the numbers this side RESOLVED — the organisation's where it has one, the deployment's otherwise — + into exactly those keys is what makes the child's budget the parent's budget. Written always, not + only when an override applies, for the posture's reason below: the child then parses a value this + process already resolved instead of repeating the resolution from text that might read differently. The one key is added because the posture is the one value the two processes must agree on that they would otherwise each read at a different MOMENT. `_pin_model_pass_posture` fixed it on this @@ -2370,11 +2493,13 @@ def _pass_child_env() -> dict[str, str]: Spelled as the canonical `true`/`false` rather than passing the operator's own text through, so the child parses a value this process has already resolved rather than repeating the resolution. """ - from execute_sql import _model_pass_disabled + from execute_sql import _model_pass_disabled, _resolve_row_cap, _resolve_timeout_s return { **os.environ, "AGAMI_GOVERNANCE_ENFORCED": "false" if _model_pass_disabled() else "true", + "AGAMI_SQL_MAX_ROWS": str(_resolve_row_cap()), + "AGAMI_SQL_TIMEOUT_S": str(_resolve_timeout_s()), } @@ -2484,14 +2609,14 @@ def _tool_execute_sql(args: dict[str, Any]) -> str: # `failed`/`timeout` naming nothing the caller can act on. Imported lazily for the same # 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. + # Resolved HERE and enforced on a child that re-resolves for itself, which only works because both + # sides reach the identical number. A request-scoped override is the one thing that could break + # that — it outranks the environment on this side of the fork, and unless it is carried across, 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. The per-organisation limit (#329) is such an override, so it is + # carried: `tool_execute_sql` pinned it before this line, this bound reads the pin, and + # `_pass_child_env` below writes the same numbers into the child's `AGAMI_SQL_*` keys. + # `_resolve_timeout_s` documents the construction, and a test drives it across a real fork. import execute_sql supervisor_timeout_s = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S @@ -3184,33 +3309,88 @@ def require_thread_id(registry: dict[str, dict[str, Any]]) -> dict[str, dict[str return out -def statement_limits() -> dict[str, int]: - """The row cap and per-statement deadline this deployment enforces, read from the same resolvers - the executor uses — so anything that shows them (the tool description, an admin screen) cannot - disagree with the bound actually applied.""" - from execute_sql import _resolve_row_cap, _resolve_timeout_s +def statement_limits(org_id: str | None = None) -> dict[str, int]: + """The row cap and per-statement deadline enforced for an organisation — ``org_id``, or the + current request's when omitted — computed the way the executor's own call computes them, so + anything that shows them (the tool description, an admin screen) cannot disagree with the bound + actually applied. + + Inside a call that already pinned its limits, and asked about no other organisation, the pin is + returned rather than a fresh resolution: that is the number this call is enforcing, even if the + provider would now answer differently.""" + import execute_sql + + pinned = execute_sql._statement_limits.get() + if org_id is None and pinned is not None: + max_rows, timeout_s = pinned + else: + max_rows, timeout_s = _effective_statement_limits(org_id) + return {"max_rows": max_rows, "timeout_s": timeout_s} + + +def statement_limit_defaults() -> dict[str, dict[str, int]]: + """What applies to an organisation with no limits of its own, and what we recommend. + + ``deployment`` is the operator's environment (``AGAMI_SQL_MAX_ROWS`` / ``AGAMI_SQL_TIMEOUT_S``); + ``recommended`` is the shipped default (1000 rows, 30 seconds). They are the same number until an + operator moves the environment, which is why an admin screen needs both: "reset to default" means + the deployment value, while "recommended" is advice that holds on any deployment.""" + import execute_sql - return {"max_rows": _resolve_row_cap(), "timeout_s": _resolve_timeout_s()} + return { + "deployment": { + "max_rows": execute_sql._row_cap_from_env(), + "timeout_s": execute_sql._timeout_s_from_env(), + }, + "recommended": { + "max_rows": execute_sql._DEFAULT_MAX_ROWS, + "timeout_s": execute_sql._DEFAULT_TIMEOUT_S, + }, + } def _execute_sql_limits_sentence() -> str: - """The limits, stated to the client before it writes SQL (#326). + """The limits, stated to the client before it writes SQL (#326), for the current organisation. The description used to name "the deployment row ceiling" and "a per-statement deadline" without either number, so a client learned them by being refused — a warehouse round trip and a retry - each time. Built when the registry is built: both resolvers read only the process environment, - which is fixed at start-up, so the number stated is the number enforced.""" + each time. The registry is built once, at import, when no organisation is known; the numbers in + it are then REPLACED per caller at list-tools time by `tool_description`, because with a + per-organisation limit (#329) the start-up numbers can be wrong for the organisation reading them.""" limits = statement_limits() return ( # "Refused", not "cancelled": some executors can only stop waiting at the bound, and the claim # the client relies on is that no answer comes back — not what happens to the work behind it. - f"Limits on this deployment: a result over {limits['max_rows']:,} rows is refused, and so is " + f"Limits in force for you: a result over {limits['max_rows']:,} rows is refused, and so is " f"a statement still running after about {limits['timeout_s']}s. Plan for both before " "running: bound a listing with ORDER BY and LIMIT, and group a breakdown more coarsely or " "filter its time range first.\n" ) +# The sentence as the registry was built with it, at import, with no organisation known. Kept by name +# so `tool_description` can find exactly this text and swap in the caller's numbers — matching the +# import-time string rather than re-rendering it means an operator changing the environment after +# start-up cannot make the search miss. +_EXECUTE_SQL_LIMITS_AT_IMPORT = _execute_sql_limits_sentence() + + +def tool_description(name: str, description: str) -> str: + """A tool's description as THIS caller should read it. + + Only `execute_sql` varies: it states the row cap and deadline, and those are per organisation + (#329), so the import-time numbers are replaced with the current request's. Called by both MCP + servers at list-tools time, inside the request whose organisation is already set. A description + without the import-time sentence — a consumer's own `execute_sql` under that name — is returned + untouched, because rewriting text we did not write would be reshaping someone else's tool. + + A client caches the list for its session, so a changed limit reaches new sessions; an existing one + meets the new numbers in the refusal, which re-resolves them per call.""" + if name != "execute_sql" or _EXECUTE_SQL_LIMITS_AT_IMPORT not in description: + return description + return description.replace(_EXECUTE_SQL_LIMITS_AT_IMPORT, _execute_sql_limits_sentence()) + + TOOLS: dict[str, dict[str, Any]] = { "list_datasources": { "handler": tool_list_datasources, @@ -3392,7 +3572,8 @@ def _execute_sql_limits_sentence() -> str: "than the deployment row ceiling (refused rather than trimmed, so a partial answer " "never arrives looking whole).\n" # The numbers behind the two limits above (#326); see `_execute_sql_limits_sentence`. - + _execute_sql_limits_sentence() + # Replaced per caller at list-tools time (#329); see `tool_description`. + + _EXECUTE_SQL_LIMITS_AT_IMPORT + # The failure channel shipped from the start and the description documented two of the # three statuses, so a client met this shape for the first time at the moment it was diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 896bb7c6..8b3f5126 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -74,7 +74,7 @@ import urllib.parse import uuid from collections.abc import Callable, Iterator -from contextvars import ContextVar, copy_context +from contextvars import ContextVar, Token, copy_context from dataclasses import asdict, dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -1514,19 +1514,39 @@ def _pin_model_pass_posture() -> bool: return value -def _resolve_row_cap() -> int: - """Effective result-row cap. `AGAMI_SQL_MAX_ROWS` is the operator-configurable DEPLOYMENT cap - (default 1000 when unset) — an operator owns their availability tradeoff and may set it higher OR - lower than 1000; it is NOT a hard 1000 ceiling. A missing/invalid/zero env value falls back to - 1000. - - The operator is the only voice here. A per-call override used to be able to lower it, and it went - with the trim (ACE-087): the one thing a caller might know better than the deployment — that it - wants MORE rows — is the thing a lowering-only override structurally could not express, and a - caller that wants 200 rows says so in the statement, where the intent is legible to everything - downstream.""" +# The effective (row cap, timeout seconds) for THIS call, when the caller's organisation has its own +# (#329). Absent means "the deployment's environment", which is every call that did not come through +# a pinning entry point. Both numbers in one value rather than two ContextVars, so a reader can never +# see one organisation's row cap beside another's deadline. +# +# This is a second, higher-precedence input to the budget, which is exactly the hazard +# `_resolve_timeout_s` used to rule out by having no such thing: a value that outranks the environment +# in the parent and is invisible to a forked child would let the supervisor bound the parent derives +# sit below the budget the child enforces. It is allowed now on the same terms `_pass_posture` was: +# it is resolved ONCE per call, before any bound is derived, and `tools._pass_child_env` writes both +# numbers into the child's `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`, so the child re-resolves the +# identical budget from its environment. The fork carries it explicitly; it is never lost across it. +# +# The values held here are already validated positive ints — the provider that supplies them lives in +# `tools`, which does the checking, so nothing here needs to re-parse them. +_statement_limits: ContextVar[tuple[int, int] | None] = ContextVar( + "_statement_limits", default=None +) + + +def _pin_statement_limits(max_rows: int, timeout_s: int) -> Token[tuple[int, int] | None]: + """Fix this call's effective budget. Returns the token the caller must reset with, in a + `finally`: the pin is request-scoped, and one left behind on a reused worker would hand the next + organisation's call this one's limits.""" + return _statement_limits.set((max_rows, timeout_s)) + + +def _row_cap_from_env() -> int: + """The DEPLOYMENT row cap: `AGAMI_SQL_MAX_ROWS`, default 1000 when unset. An operator owns their + availability tradeoff and may set it higher OR lower than 1000; it is NOT a hard 1000 ceiling. A + missing/invalid/zero env value falls back to 1000.""" raw = os.environ.get("AGAMI_SQL_MAX_ROWS", "").strip() - # `isdecimal`, not `isdigit`, for the reason `_resolve_timeout_s` gives: `isdigit` admits `²`, + # `isdecimal`, not `isdigit`, for the reason `_timeout_s_from_env` gives: `isdigit` admits `²`, # which `int()` then refuses. `tools` now resolves this while building its registry, so a raise # here would stop the module importing at all rather than failing one call. cap = int(raw) if raw.isdecimal() else _DEFAULT_MAX_ROWS @@ -1535,6 +1555,22 @@ def _resolve_row_cap() -> int: return cap +def _resolve_row_cap() -> int: + """Effective result-row cap for THIS call: the organisation's own when a call pinned one + (`_statement_limits`, #329), otherwise the deployment's (`_row_cap_from_env`). + + No per-CALL override exists, and that is unchanged. One used to be able to lower the cap, and it + went with the trim (ACE-087): the one thing a caller might know better than the deployment — that + it wants MORE rows — is the thing a lowering-only override structurally could not express, and a + caller that wants 200 rows says so in the statement, where the intent is legible to everything + downstream. What the pin carries is an administrator's setting for the organisation, resolved + before the call and not chosen by it.""" + pinned = _statement_limits.get() + if pinned is not None: + return pinned[0] + return _row_cap_from_env() + + _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 @@ -1581,18 +1617,34 @@ def _resolve_row_cap() -> int: 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. - - **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 + """Effective per-statement timeout for THIS call, in whole seconds: the organisation's own when + a call pinned one (`_statement_limits`, #329), otherwise the deployment's (`_timeout_s_from_env`). + Every bound in the ordered family — watchdog, native skew, outer bound, supervisor — reads this, + so they all derive from the one effective budget. + + **Two sources, and one budget on both sides of the fork.** The environment used to be 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. A per-organisation limit needs exactly such an override, + so the hazard is now closed by construction instead of by absence: the pin is set once per call, + before the supervisor bound is derived, and `tools._pass_child_env` writes the pinned numbers into + the child's environment, where this same resolver (with no pin of its own) reads them back. The + parent and the child therefore reach the identical number, which is the property the old rule + existed to guarantee.""" + pinned = _statement_limits.get() + if pinned is not None: + return pinned[1] + return _timeout_s_from_env() + + +def _timeout_s_from_env() -> int: + """The DEPLOYMENT per-statement timeout: `AGAMI_SQL_TIMEOUT_S`, 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. It is also what an organisation with no limit of its own gets, + and what a forked child reads the pinned budget back from. + + Unlike `_row_cap_from_env`, 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 @@ -1711,8 +1763,9 @@ def _resource_limit_refusal(exc: _ResourceLimit | None) -> Refusal: invariant that survives is one rule with one emit site, not one sentence. The budget is re-resolved rather than carried: nothing between the engine call and here can - change the environment the resolvers read, so both `_resolve_timeout_s` and `_resolve_row_cap` - return the same number the bound itself used. The configured number belongs in the detail — it + change what the resolvers read — the call's pinned organisation limits, or the environment when + none were pinned (and, in a forked child, the environment its parent wrote the pin into) — so + both `_resolve_timeout_s` and `_resolve_row_cap` return the same number the bound itself 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. """ @@ -2636,8 +2689,9 @@ def _execute_bounded( The call runs inside a copy of the CALLER's context, and what that is FOR changed when the per-call row cap went (ACE-087). It used to carry ``_max_rows_override`` to ``_resolve_row_cap`` - inside the worker; the cap is the deployment's environment now and needs no carrier. What it - still carries is the *caller's* request scope — ``tools._current_org_ctx``, the resolve-once + inside the worker; that override is gone. It now also carries the call's pinned organisation + limits (``_statement_limits``, #329), though nothing in the worker reads them today — every bound + is derived on the caller's side. What it chiefly carries is the *caller's* request scope — ``tools._current_org_ctx``, the resolve-once request cache, the actor and session on the served path — into the one place a consumer's own code runs. That is the point of the ``Executor`` seam: a pooled / per-user-RBAC executor picks its connection from exactly that context, and a new thread starts with an empty one, so dropping @@ -2742,8 +2796,9 @@ def execute_guarded( ``_load_credentials`` sits INSIDE the try deliberately, so a bad profile / missing DSN becomes a ``failed``/``dsn`` Envelope carrying its detailed message rather than escaping as an exception - the two callers would each have to translate. The row cap is the deployment's alone - (``AGAMI_SQL_MAX_ROWS``); no caller can lower it for one call.""" + the two callers would each have to translate. The row cap is the organisation's + limit when the caller pinned one (``tools.pinned_statement_limits``) and the deployment's + ``AGAMI_SQL_MAX_ROWS`` otherwise; no caller can change it for one statement.""" # Clear before anything can set it, so a detail from a PREVIOUS call in this context can never # be attributed to this one. The recorder reads it unconditionally; a stale value would put the # wrong error text on a row that succeeded. diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py index 981372ad..9a7a208f 100644 --- a/tests/test_ace038_timeout.py +++ b/tests/test_ace038_timeout.py @@ -249,6 +249,13 @@ def test_the_budget_has_exactly_one_configuration_surface(): # Written when an executor can name its statement, never read to compute a bound, and cleared at # the entry to every call beside the two above. Like them it cannot cross the fork, and like them # it does not need to — the column is null on that surface by the same construction. + # + # `_statement_limits` (#329) is the eighth, and unlike every one above it IS a budget input: an + # organisation's own row cap and deadline, outranking the environment. It is allowed through on + # `_pass_posture`'s terms, which are the only terms that answer this test's hazard: it is resolved + # once per call before any bound is derived, and `tools._pass_child_env` writes both numbers into + # the child's `AGAMI_SQL_*` keys, so the fork carries it explicitly instead of losing it. + # `tests/test_per_org_statement_limits.py` drives that across a real process boundary. context_vars -= { "_last_error_detail", "_guard_model", @@ -257,6 +264,7 @@ def test_the_budget_has_exactly_one_configuration_surface(): "_pass_posture", "_last_executing_identity", "_last_warehouse_query_id", + "_statement_limits", } assert context_vars == set(), ( "a second, higher-precedence configuration surface for the budget cannot cross the fork; " @@ -1293,7 +1301,14 @@ def _fake_run(cmd, **kwargs): # A subset rather than equality: the pinned key matches the environment whenever the environment # already spells the same posture, which is the ordinary case and the case this suite runs in. The # property being defended is that NOTHING ELSE differs. - assert differing <= {"AGAMI_GOVERNANCE_ENFORCED"}, differing + # + # #329 adds the two budget keys, and they are the opposite of a disturbance: they are written so + # the child reaches the budget this side resolved (an organisation's own limit, when it has one). + # They may differ from `os.environ` only in SPELLING an unset or unusable value as the number it + # resolves to — asserted by value, so a child handed any other budget still fails here. + assert differing <= {"AGAMI_GOVERNANCE_ENFORCED", "AGAMI_SQL_MAX_ROWS", "AGAMI_SQL_TIMEOUT_S"}, differing + assert child_env["AGAMI_SQL_TIMEOUT_S"] == "300" + assert child_env["AGAMI_SQL_MAX_ROWS"] == str(execute_sql._row_cap_from_env()) def test_the_supervisors_verdict_is_unchanged(warehouse, monkeypatch): diff --git a/tests/test_per_org_statement_limits.py b/tests/test_per_org_statement_limits.py new file mode 100644 index 00000000..75ba0d41 --- /dev/null +++ b/tests/test_per_org_statement_limits.py @@ -0,0 +1,391 @@ +"""An organisation's own row cap and statement deadline (#329). + +Core stores no such setting. A consumer registers a provider, `(org_id) -> {"max_rows", "timeout_s"}`, +and core resolves it ONCE per call, holds the answer for the whole call, and writes the same numbers +into a forked child's environment — so the ordered bound family (watchdog < native < outer < +supervisor) stays one budget on both sides of the fork. Anything missing or unusable falls back to +the deployment's `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`, and there is no ceiling. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("pydantic") + +import execute_sql # noqa: E402 +import tools # noqa: E402 + +PKG_SRC = Path(__file__).resolve().parent.parent / "packages" / "agami-core" / "src" + +_LIMITS = { + "acme": {"max_rows": 5000, "timeout_s": 90}, + "globex": {"max_rows": 20, "timeout_s": 4}, +} + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + monkeypatch.delenv("AGAMI_SQL_MAX_ROWS", raising=False) + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) + tools.set_statement_limits_provider(None) + yield + tools.set_statement_limits_provider(None) + + +def _in_org(monkeypatch, org_id: str) -> None: + # Patched rather than set on `_current_org_ctx`, so monkeypatch undoes it and no later test starts + # inside this organisation. + monkeypatch.setattr(tools, "_current_org_id", lambda: org_id) + + +def _by_org(org_id: str): + return _LIMITS.get(org_id) + + +# ---------------------------------------------------------------------------------------------- +# The pin: in-process resolvers see the organisation's numbers, and only for the call +# ---------------------------------------------------------------------------------------------- + + +def test_the_override_reaches_both_in_process_resolvers(): + tools.set_statement_limits_provider(_by_org) + + with tools.pinned_statement_limits("acme") as pinned: + assert pinned == (5000, 90) + assert execute_sql._resolve_row_cap() == 5000 + assert execute_sql._resolve_timeout_s() == 90 + + assert execute_sql._resolve_row_cap() == execute_sql._DEFAULT_MAX_ROWS + assert execute_sql._resolve_timeout_s() == execute_sql._DEFAULT_TIMEOUT_S + + +def test_two_organisations_in_sequence_do_not_leak_into_each_other(): + """One worker thread serves one organisation after another; a pin that outlived its call would + hand the next organisation these limits.""" + tools.set_statement_limits_provider(_by_org) + + with tools.pinned_statement_limits("acme"): + assert execute_sql._resolve_timeout_s() == 90 + assert execute_sql._statement_limits.get() is None + + with tools.pinned_statement_limits("globex"): + assert (execute_sql._resolve_row_cap(), execute_sql._resolve_timeout_s()) == (20, 4) + assert execute_sql._statement_limits.get() is None + + with tools.pinned_statement_limits("initech"): # no row for this one + assert execute_sql._resolve_timeout_s() == execute_sql._DEFAULT_TIMEOUT_S + + +def test_the_pin_is_released_when_the_call_raises(): + tools.set_statement_limits_provider(_by_org) + + with pytest.raises(RuntimeError), tools.pinned_statement_limits("acme"): + raise RuntimeError("boom") + + assert execute_sql._statement_limits.get() is None + + +def test_tool_execute_sql_holds_the_callers_limits_for_the_whole_call(monkeypatch): + tools.set_statement_limits_provider(_by_org) + _in_org(monkeypatch, "globex") + seen: dict = {} + + def _body(args): + seen["limits"] = (execute_sql._resolve_row_cap(), execute_sql._resolve_timeout_s()) + return "{}" + + monkeypatch.setattr(tools, "_tool_execute_sql", _body) + tools.tool_execute_sql({"sql": "SELECT 1"}) + + assert seen["limits"] == (20, 4) + assert execute_sql._statement_limits.get() is None + + +def test_the_provider_is_asked_once_per_call(monkeypatch): + """Once, not once per reader: a provider whose answer changed mid-call must not be able to give + the supervisor bound and the child environment different numbers.""" + calls: list[str] = [] + + def _counting(org_id): + calls.append(org_id) + return {"max_rows": 5000 + len(calls), "timeout_s": 90 + len(calls)} + + tools.set_statement_limits_provider(_counting) + _in_org(monkeypatch, "acme") + seen: dict = {} + + def _body(args): + first = tools._pass_child_env() + seen["child"] = (first["AGAMI_SQL_MAX_ROWS"], first["AGAMI_SQL_TIMEOUT_S"]) + seen["supervisor"] = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S + seen["reported"] = tools.statement_limits() + return "{}" + + monkeypatch.setattr(tools, "_tool_execute_sql", _body) + tools.tool_execute_sql({"sql": "SELECT 1"}) + + assert calls == ["acme"] + assert seen["child"] == ("5001", "91") + assert seen["supervisor"] == 91 + execute_sql._SUPERVISOR_SKEW_S + assert seen["reported"] == {"max_rows": 5001, "timeout_s": 91} + + +# ---------------------------------------------------------------------------------------------- +# The fork: the child is handed, and really resolves, the same budget +# ---------------------------------------------------------------------------------------------- + + +def test_the_fork_path_hands_the_child_the_organisations_budget(monkeypatch): + tools.set_statement_limits_provider(_by_org) + _in_org(monkeypatch, "acme") + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + monkeypatch.setattr(tools, "_INJECTED_EXECUTOR", None) + captured: dict = {} + + class _Stop(Exception): + pass + + def _fake_run(cmd, **kwargs): + captured.update(kwargs) + raise _Stop # the budget is decided by now; nothing after the fork is under test + + monkeypatch.setattr(tools.subprocess, "run", _fake_run) + monkeypatch.setattr(tools, "_resolve_call_datasource", lambda args: "demo") + + with pytest.raises(_Stop): + tools.tool_execute_sql({"sql": "SELECT 1"}) + + assert captured["env"]["AGAMI_SQL_MAX_ROWS"] == "5000" + assert captured["env"]["AGAMI_SQL_TIMEOUT_S"] == "90" + assert captured["timeout"] == 90 + execute_sql._SUPERVISOR_SKEW_S + + +def test_a_real_child_resolves_the_budget_the_parent_bounded(): + """Driven across the actual process boundary. The child has no provider and no pin — only the + environment `_pass_child_env` built — and must reach the parent's number, or the supervisor + stops waiting inside the budget the child is still enforcing.""" + tools.set_statement_limits_provider(_by_org) + + with tools.pinned_statement_limits("acme"): + env = tools._pass_child_env() + 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_row_cap(), execute_sql._resolve_timeout_s())", + ], + capture_output=True, + text=True, + timeout=60, + env={**env, "PYTHONPATH": str(PKG_SRC)}, + ) + assert child.returncode == 0, child.stderr + child_rows, child_timeout = (int(x) for x in child.stdout.split()) + + assert (child_rows, child_timeout) == (5000, 90) + assert parent_bound > child_timeout + + +def test_without_a_provider_the_child_env_spells_the_deployment_values(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_MAX_ROWS", "250") + + with tools.pinned_statement_limits("acme"): + env = tools._pass_child_env() + + assert env["AGAMI_SQL_MAX_ROWS"] == "250" + assert env["AGAMI_SQL_TIMEOUT_S"] == str(execute_sql._DEFAULT_TIMEOUT_S) + + +# ---------------------------------------------------------------------------------------------- +# Fallbacks: never raise, fall back to the environment +# ---------------------------------------------------------------------------------------------- + + +def test_no_provider_means_the_deployment_values(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_MAX_ROWS", "300") + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "12") + + with tools.pinned_statement_limits("acme") as pinned: + assert pinned == (300, 12) + + +@pytest.mark.parametrize( + "answer", + [None, {}, {"max_rows": None, "timeout_s": None}], +) +def test_an_organisation_without_its_own_limits_gets_the_deployment_values(monkeypatch, answer): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "12") + tools.set_statement_limits_provider(lambda org_id: answer) + + with tools.pinned_statement_limits("acme") as pinned: + assert pinned == (execute_sql._DEFAULT_MAX_ROWS, 12) + + +def test_one_key_may_be_overridden_alone(): + tools.set_statement_limits_provider(lambda org_id: {"timeout_s": 600}) + + assert tools.statement_limits("acme") == { + "max_rows": execute_sql._DEFAULT_MAX_ROWS, + "timeout_s": 600, + } + + +@pytest.mark.parametrize("bad", [0, -1, "50", 2.5, True, [10]]) +def test_an_unusable_value_falls_back_with_a_warning(monkeypatch, caplog, bad): + monkeypatch.setenv("AGAMI_SQL_MAX_ROWS", "300") + tools.set_statement_limits_provider(lambda org_id: {"max_rows": bad, "timeout_s": 45}) + + with caplog.at_level(logging.WARNING, logger=tools._LOG.name): + limits = tools.statement_limits("acme") + + assert limits == {"max_rows": 300, "timeout_s": 45} # the good key still applies + assert any("max_rows" in r.getMessage() for r in caplog.records) + + +def test_there_is_no_ceiling(): + """The owner's decision: an administrator may set any positive number. Stated as a test so a + clamp added later is a deliberate change rather than a quiet one.""" + tools.set_statement_limits_provider( + lambda org_id: {"max_rows": 10_000_000, "timeout_s": 86_400} + ) + + assert tools.statement_limits("acme") == {"max_rows": 10_000_000, "timeout_s": 86_400} + + +def test_a_provider_that_raises_falls_back_and_does_not_fail_the_call(caplog): + def _broken(org_id): + raise ConnectionError("settings store unavailable") + + tools.set_statement_limits_provider(_broken) + + with caplog.at_level(logging.WARNING, logger=tools._LOG.name): + with tools.pinned_statement_limits("acme") as pinned: + assert pinned == (execute_sql._DEFAULT_MAX_ROWS, execute_sql._DEFAULT_TIMEOUT_S) + + assert any("failed" in r.getMessage() for r in caplog.records) + + +def test_a_provider_that_returns_a_non_mapping_falls_back(caplog): + tools.set_statement_limits_provider(lambda org_id: (5000, 90)) + + with caplog.at_level(logging.WARNING, logger=tools._LOG.name): + assert tools.statement_limits("acme")["timeout_s"] == execute_sql._DEFAULT_TIMEOUT_S + + assert any("not a mapping" in r.getMessage() for r in caplog.records) + + +def test_a_non_callable_provider_is_refused_at_registration(): + with pytest.raises(TypeError): + tools.set_statement_limits_provider({"max_rows": 10}) # type: ignore[arg-type] + + +# ---------------------------------------------------------------------------------------------- +# What is reported: effective limits, deployment defaults, recommended values +# ---------------------------------------------------------------------------------------------- + + +def test_statement_limits_reports_the_current_organisations_limits(monkeypatch): + tools.set_statement_limits_provider(_by_org) + _in_org(monkeypatch, "acme") + + assert tools.statement_limits() == {"max_rows": 5000, "timeout_s": 90} + assert tools.statement_limits("globex") == {"max_rows": 20, "timeout_s": 4} + + +def test_defaults_report_the_deployment_and_the_recommendation_separately(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_MAX_ROWS", "2500") + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") + tools.set_statement_limits_provider(_by_org) + + assert tools.statement_limit_defaults() == { + "deployment": {"max_rows": 2500, "timeout_s": 45}, + "recommended": {"max_rows": 1000, "timeout_s": 30}, + } + + +# ---------------------------------------------------------------------------------------------- +# What the client is told matches what is enforced for its organisation +# ---------------------------------------------------------------------------------------------- + + +def test_the_description_states_the_callers_numbers(monkeypatch): + tools.set_statement_limits_provider(_by_org) + _in_org(monkeypatch, "acme") + + described = tools.tool_description("execute_sql", tools.TOOLS["execute_sql"]["description"]) + + assert "5,000 rows" in described and "90s" in described + assert tools._execute_sql_limits_sentence() in described + assert "max_rows" not in described # ACE-087: the removed argument is never advertised + + +def test_only_our_execute_sql_description_is_rewritten(monkeypatch): + tools.set_statement_limits_provider(_by_org) + _in_org(monkeypatch, "acme") + other = tools.TOOLS["list_datasources"]["description"] + + assert tools.tool_description("list_datasources", other) == other + assert tools.tool_description("execute_sql", "a consumer's own tool") == "a consumer's own tool" + + +def test_the_http_server_lists_the_callers_numbers(monkeypatch): + pytest.importorskip("mcp") + import mcp.types as mt + import mcp_http + + tools.set_statement_limits_provider(_by_org) + server = mcp_http.build_server() + handler = server.request_handlers[mt.ListToolsRequest] + + async def _list_as(org_id: str) -> str: + token = tools._current_org_ctx.set(org_id) + try: + result = await handler(mt.ListToolsRequest(method="tools/list")) + finally: + tools._current_org_ctx.reset(token) + listed = {t.name: t.description for t in result.root.tools} + return listed["execute_sql"] + + assert "5,000 rows" in asyncio.run(_list_as("acme")) + assert "20 rows" in asyncio.run(_list_as("globex")) + + +def test_the_refusal_names_the_organisations_row_cap(): + """An existing session keeps the description it listed; the refusal re-resolves per call, so it + is where a changed limit is met first.""" + tools.set_statement_limits_provider(_by_org) + + with tools.pinned_statement_limits("globex"): + refusal = execute_sql._resource_limit_refusal(None) + + assert "20-row limit" in refusal.detail + + +def test_create_app_registers_the_adapters_provider(monkeypatch): + pytest.importorskip("mcp") + import dataclasses + + import mcp_http + + monkeypatch.setenv("PUBLIC_BASE_URL", "http://localhost:8000") + adapters = dataclasses.replace(mcp_http.default_adapters(), statement_limits=_by_org) + + mcp_http.create_app(adapters=adapters) + try: + assert tools._STATEMENT_LIMITS_PROVIDER is _by_org + finally: + tools.set_injected_executor(None) + + mcp_http.create_app() + assert tools._STATEMENT_LIMITS_PROVIDER is None + assert os.environ.get("AGAMI_SQL_MAX_ROWS") is None # registration touches no environment From b1a3331fff132725669e3d3e70bb56d80cdf535d Mon Sep 17 00:00:00 2001 From: Ashwin Ramachandran Date: Mon, 14 Sep 2026 15:50:02 -0700 Subject: [PATCH 2/3] Review fixes: refuse unarmable timeouts, pin limits in evaluation runs - A timeout whose supervisor bound reaches threading.TIMEOUT_MAX is treated as unusable (provider and AGAMI_SQL_TIMEOUT_S alike) rather than capped: at that size Timer/join raise OverflowError, and the outer bound raised it after its worker started and before the abandonment was counted, so the abandoned-worker cap stopped bounding anything. - golden_run scores each case under tools.pinned_statement_limits(org), imported lazily. - list_tools skips the thread hop when no provider is registered. - statement_limits(org_id=...) inside a call documented as asking afresh. - Tests: assertions that can fail; the outer bound and watchdog fire on the pinned timeout; the abandoned-worker cap holds under an unrepresentable one. Refs #329 Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 5 +- packages/agami-core/src/execute_sql.py | 16 +- packages/agami-core/src/mcp_http.py | 4 + .../src/semantic_model/golden_run.py | 39 ++++- packages/agami-core/src/tools.py | 27 ++- plugins/agami/lib/execute_sql.py | 16 +- tests/test_golden_run.py | 32 ++++ tests/test_per_org_statement_limits.py | 154 +++++++++++++++++- 8 files changed, 272 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bcfa463..8cb15ec7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,7 +19,10 @@ below corresponds to one such version. through `Adapters.statement_limits` or `tools.set_statement_limits_provider`. A missing, `None` or unusable value (not a positive whole number, or a provider that raises) falls back to `AGAMI_SQL_MAX_ROWS` / `AGAMI_SQL_TIMEOUT_S`, which stay the deployment default, with a warning in - the log. There is no ceiling. + the log. There is no ceiling. A time limit too large for the platform to arm a timer on (at or + above Python's `threading.TIMEOUT_MAX`, counting the supervisor's 60-second slack) is treated as + unusable, from the provider and from `AGAMI_SQL_TIMEOUT_S` alike. + - An evaluation run scores both statements of each case under the named organisation's limits. - The limits are resolved once per `execute_sql` call and held for the whole call, and the forked child is handed the same two numbers in its environment, so the watchdog, the native bound, the outer bound and the supervisor still derive from one budget on both sides of the fork. diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 8b3f5126..92433c46 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -1638,6 +1638,19 @@ def _resolve_timeout_s() -> int: return _timeout_s_from_env() +def _timeout_is_representable(timeout_s: int) -> bool: + """Whether every bound derived from `timeout_s` can actually be armed. + + Not a ceiling — there is deliberately none (#329) — but a statement of what the platform can + express. `threading.Timer` (the watchdog), `Thread.join` (the outer bound) and the supervisor's + wait all refuse a timeout at or above `threading.TIMEOUT_MAX` with an `OverflowError`. The outer + bound raises it AFTER its worker has started and BEFORE the abandonment is counted, so an + unrepresentable budget would not merely fail one call: it would leave `_MAX_ABANDONED_WORKERS` + bounding nothing. The largest derived bound is the supervisor's, so that is the one checked, and a + value failing it is treated as unusable like any other and falls back to the deployment's.""" + return timeout_s + _SUPERVISOR_SKEW_S < threading.TIMEOUT_MAX + + def _timeout_s_from_env() -> int: """The DEPLOYMENT per-statement timeout: `AGAMI_SQL_TIMEOUT_S`, default 30 when unset. An operator owns their availability tradeoff and may set it higher OR lower than 30. A missing or non-positive @@ -1657,7 +1670,8 @@ def _timeout_s_from_env() -> int: # 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 + usable = written is not None and written > 0 and _timeout_is_representable(written) + timeout_s = written if usable else _DEFAULT_TIMEOUT_S if raw and timeout_s != written: _LOG.warning( "AGAMI_SQL_TIMEOUT_S=%r is not a usable whole number of seconds; falling back to %ds.", diff --git a/packages/agami-core/src/mcp_http.py b/packages/agami-core/src/mcp_http.py index a5d83b1d..eca322b7 100644 --- a/packages/agami-core/src/mcp_http.py +++ b/packages/agami-core/src/mcp_http.py @@ -51,6 +51,7 @@ TOOLS, _current_org_ctx, bootstrap_paths, + has_statement_limits_provider, record_tool_call, require_thread_id, reset_typed_outcome, @@ -503,6 +504,9 @@ async def _list_tools() -> list: # Off the loop, for ACE-048's reason: the limits provider is the consumer's code and usually a # database read, and on the loop one slow read would stall every in-flight request. # `run_blocking` copies the request context, so the organisation is still set in the worker. + # Without a provider nothing here blocks, so the hop would be a thread per listing for nothing. + if not has_statement_limits_provider(): + return _listed() return await run_blocking(_listed) @server.call_tool() diff --git a/packages/agami-core/src/semantic_model/golden_run.py b/packages/agami-core/src/semantic_model/golden_run.py index 1499ba72..1a396b2e 100644 --- a/packages/agami-core/src/semantic_model/golden_run.py +++ b/packages/agami-core/src/semantic_model/golden_run.py @@ -44,6 +44,7 @@ from __future__ import annotations +import contextlib import json import os import shutil @@ -303,6 +304,21 @@ def run_golden_dataset( ) +def _org_statement_limits(org: str) -> contextlib.AbstractContextManager: + """`tools.pinned_statement_limits(org)`, imported at the call rather than at module load. + + The provider registry lives in `tools`, the serving layer, and this module sits below it: `tools` + reaches `semantic_model` only lazily, so a lazy import here closes no cycle, but a module-level one + would make loading the evaluation runner load the whole tool registry. Where `tools` cannot be + imported at all, no provider can have been registered either, so the deployment's limits — what + an unpinned call already enforces — are the right answer and not a degraded one.""" + try: + from tools import pinned_statement_limits + except ImportError: + return contextlib.nullcontext() + return pinned_statement_limits(org) + + def _run_item( item: GoldenItem, generated: GeneratedSql, @@ -325,15 +341,20 @@ def _run_item( ) golden_sql = item.expected.sql or "" - score = _score( - item, - generated.sql, - golden_sql, - profile=profile, - org=org, - executor=executor, - dialect=dialect, - ) + # Scored under THIS org's statement limits (#329), not the deployment's: an evaluation that runs + # against looser or tighter limits than the org's own questions do would score a statement the + # org could never have run, or refuse one it could. Both statements of the item share one pin, so + # the answer key and the generated statement are held to the same budget. + with _org_statement_limits(org): + score = _score( + item, + generated.sql, + golden_sql, + profile=profile, + org=org, + executor=executor, + dialect=dialect, + ) # After the score, and on EVERY item that produced a statement — an answer key is not the # condition. The diff is what turns "the rows disagree" into a reason, and one of its two gates # reads the generated statement alone: `must_filter` is the DATASET's requirement rather than a diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index c781ba15..5603d312 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -1929,11 +1929,18 @@ def _provider_limit(org_id: str, key: str, value: Any) -> int | None: organisation's bad row must cost that organisation its override, not everybody their query.""" if value is None: return None - if isinstance(value, int) and not isinstance(value, bool) and value > 0: + from execute_sql import _timeout_is_representable + + usable = isinstance(value, int) and not isinstance(value, bool) and value > 0 + # No ceiling, but a timeout the platform cannot arm is not a setting: see + # `execute_sql._timeout_is_representable` for how one would disable the abandoned-worker cap. + if usable and key == "timeout_s" and not _timeout_is_representable(value): + usable = False + if usable: return value _LOG.warning( - "statement limits provider returned %s=%r for org %s, which is not a positive whole number; " - "using the deployment value.", + "statement limits provider returned %s=%r for org %s, which is not a usable positive whole " + "number; using the deployment value.", key, value, org_id, @@ -1981,6 +1988,12 @@ def _effective_statement_limits(org_id: str | None) -> tuple[int, int]: return limits["max_rows"], limits["timeout_s"] +def has_statement_limits_provider() -> bool: + """Whether a provider is registered — so a caller can skip work that only a provider makes + necessary (the HTTP server's thread hop when listing tools).""" + return _STATEMENT_LIMITS_PROVIDER is not None + + @contextlib.contextmanager def pinned_statement_limits(org_id: str | None = None) -> Iterator[tuple[int, int]]: """Resolve an organisation's effective limits ONCE and hold them for the enclosed block. @@ -3315,9 +3328,11 @@ def statement_limits(org_id: str | None = None) -> dict[str, int]: anything that shows them (the tool description, an admin screen) cannot disagree with the bound actually applied. - Inside a call that already pinned its limits, and asked about no other organisation, the pin is - returned rather than a fresh resolution: that is the number this call is enforcing, even if the - provider would now answer differently.""" + Inside a call that already pinned its limits, and with ``org_id`` omitted, the pin is returned + rather than a fresh resolution: that is the number this call is enforcing, even if the provider + would now answer differently. Passing an ``org_id`` — even the current request's own — always asks + the provider afresh, because the pin records numbers and not whose they are; a caller that wants + what this call enforces omits the argument.""" import execute_sql pinned = execute_sql._statement_limits.get() diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 8b3f5126..92433c46 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -1638,6 +1638,19 @@ def _resolve_timeout_s() -> int: return _timeout_s_from_env() +def _timeout_is_representable(timeout_s: int) -> bool: + """Whether every bound derived from `timeout_s` can actually be armed. + + Not a ceiling — there is deliberately none (#329) — but a statement of what the platform can + express. `threading.Timer` (the watchdog), `Thread.join` (the outer bound) and the supervisor's + wait all refuse a timeout at or above `threading.TIMEOUT_MAX` with an `OverflowError`. The outer + bound raises it AFTER its worker has started and BEFORE the abandonment is counted, so an + unrepresentable budget would not merely fail one call: it would leave `_MAX_ABANDONED_WORKERS` + bounding nothing. The largest derived bound is the supervisor's, so that is the one checked, and a + value failing it is treated as unusable like any other and falls back to the deployment's.""" + return timeout_s + _SUPERVISOR_SKEW_S < threading.TIMEOUT_MAX + + def _timeout_s_from_env() -> int: """The DEPLOYMENT per-statement timeout: `AGAMI_SQL_TIMEOUT_S`, default 30 when unset. An operator owns their availability tradeoff and may set it higher OR lower than 30. A missing or non-positive @@ -1657,7 +1670,8 @@ def _timeout_s_from_env() -> int: # 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 + usable = written is not None and written > 0 and _timeout_is_representable(written) + timeout_s = written if usable else _DEFAULT_TIMEOUT_S if raw and timeout_s != written: _LOG.warning( "AGAMI_SQL_TIMEOUT_S=%r is not a usable whole number of seconds; falling back to %ds.", diff --git a/tests/test_golden_run.py b/tests/test_golden_run.py index 30bf8d19..53bf46e7 100644 --- a/tests/test_golden_run.py +++ b/tests/test_golden_run.py @@ -1035,3 +1035,35 @@ def _missing(*args, **kwargs): assert generated.sql == "" and generated.error == gr._GENERATION_UNAVAILABLE + + +# --- the org's statement limits (#329) --------------------------------------------------------- + + +def test_both_statements_are_scored_under_the_orgs_own_statement_limits(chokepoint): + """A hosted evaluation names its org; the statements it runs must meet that org's limits, not + the deployment's, or the run scores a statement the org's own questions could never have run.""" + import tools + + seen: list[tuple[int, int]] = [] + + class _LimitsSpy(_SpyExecutor): + def execute(self, vetted_sql, creds, *, profile): + seen.append((execute_sql._resolve_row_cap(), execute_sql._resolve_timeout_s())) + return super().execute(vetted_sql, creds, profile=profile) + + asked: list[str] = [] + + def _provider(org_id): + asked.append(org_id) + return {"max_rows": 4321, "timeout_s": 77} + + tools.set_statement_limits_provider(_provider) + try: + _run(_dataset(), _StubGenerator(), _LimitsSpy()) + finally: + tools.set_statement_limits_provider(None) + + assert asked and set(asked) == {ORG} + assert seen == [(4321, 77), (4321, 77)] # the answer key and the generated statement + assert execute_sql._statement_limits.get() is None diff --git a/tests/test_per_org_statement_limits.py b/tests/test_per_org_statement_limits.py index 75ba0d41..c38e6696 100644 --- a/tests/test_per_org_statement_limits.py +++ b/tests/test_per_org_statement_limits.py @@ -11,9 +11,10 @@ import asyncio import logging -import os import subprocess import sys +import threading +import time from pathlib import Path import pytest @@ -193,7 +194,9 @@ def test_a_real_child_resolves_the_budget_the_parent_bounded(): child_rows, child_timeout = (int(x) for x in child.stdout.split()) assert (child_rows, child_timeout) == (5000, 90) - assert parent_bound > child_timeout + # The parent's supervisor bound is derived from the SAME number the child enforces — not merely + # larger, which the fixed skew alone would make true of any budget the child might have reached. + assert parent_bound - execute_sql._SUPERVISOR_SKEW_S == child_timeout def test_without_a_provider_the_child_env_spells_the_deployment_values(monkeypatch): @@ -388,4 +391,149 @@ def test_create_app_registers_the_adapters_provider(monkeypatch): mcp_http.create_app() assert tools._STATEMENT_LIMITS_PROVIDER is None - assert os.environ.get("AGAMI_SQL_MAX_ROWS") is None # registration touches no environment + # And behaviourally: acme's own limits no longer apply once an app without a provider is built. + assert tools.statement_limits("acme") == { + "max_rows": execute_sql._DEFAULT_MAX_ROWS, + "timeout_s": execute_sql._DEFAULT_TIMEOUT_S, + } + + +def test_listing_tools_without_a_provider_does_not_hop_threads(monkeypatch): + pytest.importorskip("mcp") + import mcp.types as mt + import mcp_http + + async def _no_hop(*args, **kwargs): + raise AssertionError("listed on a worker thread with no provider registered") + + monkeypatch.setattr(mcp_http, "run_blocking", _no_hop) + handler = mcp_http.build_server().request_handlers[mt.ListToolsRequest] + + result = asyncio.run(handler(mt.ListToolsRequest(method="tools/list"))) + + assert "execute_sql" in {t.name for t in result.root.tools} + + +def test_naming_an_org_inside_a_call_asks_the_provider_again(): + """The pin records numbers, not whose they are, so only the argument-less form returns it.""" + calls: list[str] = [] + + def _counting(org_id): + calls.append(org_id) + return {"timeout_s": 90 + len(calls)} + + tools.set_statement_limits_provider(_counting) + + with tools.pinned_statement_limits("acme"): + assert tools.statement_limits()["timeout_s"] == 91 + assert tools.statement_limits("acme")["timeout_s"] == 92 + + assert calls == ["acme", "acme"] + + +# ---------------------------------------------------------------------------------------------- +# The bounds themselves run on the pinned budget, and every budget is one they can arm +# ---------------------------------------------------------------------------------------------- + + +class _Blocking: + """An executor that does not return until released — the shape the outer bound exists for.""" + + def __init__(self) -> None: + self.release = threading.Event() + + def execute(self, vetted_sql, creds, *, profile): + self.release.wait(30) + return execute_sql.ExecResult(columns=["c"], rows=[(1,)], truncated=False) + + +def _drain_abandoned(deadline_s: float = 5) -> None: + deadline = time.monotonic() + deadline_s + while execute_sql._abandoned_workers and time.monotonic() < deadline: + time.sleep(0.01) + + +def test_the_outer_bound_fires_on_the_pinned_timeout_not_the_deployments(monkeypatch): + """The resolvers returning the pin is not the claim; the bound waiting that long is. With the + deployment at 30s and acme at 1s, the outer bound must stop waiting after about one second.""" + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + tools.set_statement_limits_provider(lambda org_id: {"timeout_s": 1}) + blocking = _Blocking() + + started = time.monotonic() + try: + with tools.pinned_statement_limits("acme"), pytest.raises(execute_sql._OuterBoundExpired): + execute_sql._execute_bounded(blocking, "SELECT 1", {}, profile="demo") + elapsed = time.monotonic() - started + finally: + blocking.release.set() + _drain_abandoned() + + assert elapsed < 10, ( + f"the outer bound waited {elapsed:.1f}s — the deployment's budget, not acme's" + ) + + +def test_the_watchdog_fires_on_the_pinned_timeout(monkeypatch): + tools.set_statement_limits_provider(lambda org_id: {"timeout_s": 1}) + cancelled = threading.Event() + + started = time.monotonic() + with tools.pinned_statement_limits("acme"): + with execute_sql._deadline(cancelled.set, execute_sql._resolve_timeout_s()) as fired: + cancelled.wait(10) + elapsed = time.monotonic() - started + + assert fired.is_set() and cancelled.is_set() + assert elapsed < 10 + + +_UNREPRESENTABLE = int(threading.TIMEOUT_MAX) +_LARGEST_ARMABLE = int(threading.TIMEOUT_MAX) - execute_sql._SUPERVISOR_SKEW_S - 1 + + +def test_a_provider_timeout_the_platform_cannot_arm_falls_back(caplog): + tools.set_statement_limits_provider(lambda org_id: {"timeout_s": _UNREPRESENTABLE}) + + with caplog.at_level(logging.WARNING, logger=tools._LOG.name): + assert tools.statement_limits("acme")["timeout_s"] == execute_sql._DEFAULT_TIMEOUT_S + + assert any("timeout_s" in r.getMessage() for r in caplog.records) + + +def test_the_largest_armable_timeout_is_still_accepted(): + """Representability, not a ceiling: one second below the edge is the administrator's to set.""" + tools.set_statement_limits_provider(lambda org_id: {"timeout_s": _LARGEST_ARMABLE}) + + assert tools.statement_limits("acme")["timeout_s"] == _LARGEST_ARMABLE + + +@pytest.mark.parametrize("raw", [str(_UNREPRESENTABLE), "99999999999999999999"]) +def test_an_environment_timeout_the_platform_cannot_arm_falls_back(monkeypatch, caplog, raw): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + assert execute_sql._timeout_s_from_env() == execute_sql._DEFAULT_TIMEOUT_S + + assert any("AGAMI_SQL_TIMEOUT_S" in r.getMessage() for r in caplog.records) + + +def test_an_unrepresentable_timeout_cannot_disable_the_abandoned_worker_cap(monkeypatch): + """The reported failure: `join(huge)` raised OverflowError after the worker started and before + the abandonment was counted, so the cap stopped bounding anything. The fallback budget is shrunk + to one second so the abandonment happens inside the test; the cap is shrunk to one so a single + abandonment reaches it.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_UNREPRESENTABLE)) + monkeypatch.setattr(execute_sql, "_DEFAULT_TIMEOUT_S", 1) + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + monkeypatch.setattr(execute_sql, "_MAX_ABANDONED_WORKERS", execute_sql._abandoned_workers + 1) + blocking = _Blocking() + + try: + with pytest.raises(execute_sql._OuterBoundExpired): # not OverflowError + execute_sql._execute_bounded(blocking, "SELECT 1", {}, profile="demo") + with pytest.raises(execute_sql._ExecutorSaturated): + execute_sql._execute_bounded(blocking, "SELECT 1", {}, profile="demo") + finally: + blocking.release.set() + _drain_abandoned() From b5e9b5a8ea7086d8ab856b92a4ab2eec5071895a Mon Sep 17 00:00:00 2001 From: Ashwin Ramachandran Date: Mon, 14 Sep 2026 16:10:07 -0700 Subject: [PATCH 3/3] CHANGELOG: one Added section under Unreleased after merging main Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b81b8f65..cb9c003d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,12 @@ below corresponds to one such version. - The `execute_sql` description states the caller's organisation's numbers, built when tools are listed rather than once at start-up. A client keeps the list for its session, so a changed limit reaches new sessions; an existing session meets it in the refusal, which names the number per call. +- **`sm set-description`, and onboarding asks for a datasource description** (#327). The one line + `list_datasources` shows an agent to route a question by could only be hand-edited into + `datasource.yaml`. `sm set-description --description "…"` writes it (validated, committed), + `agami-connect` asks for it on every onboard — with an option to generate it from the enriched + model, as it does for the database narrative — and `model_deploy` warns when a datasource is + deployed without one. ### Fixed @@ -66,15 +72,6 @@ below corresponds to one such version. semantic-model pass, which is off by default on a server (see `SECURITY.md`), so there the hint has nothing to rewrite. -### Added - -- **`sm set-description`, and onboarding asks for a datasource description** (#327). The one line - `list_datasources` shows an agent to route a question by could only be hand-edited into - `datasource.yaml`. `sm set-description --description "…"` writes it (validated, committed), - `agami-connect` asks for it on every onboard — with an option to generate it from the enriched - model, as it does for the database narrative — and `model_deploy` warns when a datasource is - deployed without one. - ## [0.8.6] — 2026-09-14 ### Added