Skip to content

Let an organisation have its own row cap and statement time limit (engine seam) - #334

Merged
ashwin-agami merged 4 commits into
mainfrom
per-org-statement-limits
Sep 14, 2026
Merged

ashwin-agami merged 4 commits into
mainfrom
per-org-statement-limits

Conversation

@ashwin-agami

@ashwin-agami ashwin-agami commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

Refs #329 — the engine half. The admin screen, its storage and its audit live in the private product and are not in this PR.

What it does

An organisation can have its own result row cap and per-statement time limit. Core stores nothing: it exposes a seam, asks it once per call, and holds the answer identical on both sides of the fork. AGAMI_SQL_MAX_ROWS / AGAMI_SQL_TIMEOUT_S stay the deployment default for any organisation without its own.

The issue proposed a deployment ceiling; the owner decided against one (see the availability note below), so this PR clamps nothing.

Public interface

In tools:

  • set_statement_limits_provider(provider: Callable[[str], Mapping[str, Any] | None] | None) -> None registers or clears the provider. provider(org_id) returns {"max_rows": int | None, "timeout_s": int | None} or None. A missing key or a None means use the deployment value. A non-callable raises TypeError at registration.
  • ports.Adapters.statement_limits is the same callable carried on the adapters. mcp_http.create_app registers it unconditionally, the way it registers executor, so an app built without one clears any earlier registration.
  • statement_limits(org_id: str | None = None) -> dict[str, int] returns the limits in force for the current request's organisation, or for a named one. Inside a call it returns that call's pinned numbers.
  • statement_limit_defaults() -> dict[str, dict[str, int]] returns {"deployment": {max_rows, timeout_s}, "recommended": {"max_rows": 1000, "timeout_s": 30}}.
  • pinned_statement_limits(org_id: str | None = None) is a context manager that resolves once and holds the limits for the block. tool_execute_sql opens it around every call. A direct caller of execute_sql.execute_guarded (an evaluation run, for example) opens it to get the organisation's limits rather than the deployment's.
  • tool_description(name: str, description: str) -> str returns a description as the current caller should read it. Both MCP servers call it at list-tools time.

Validation: a value that is not a positive int (0, -1, "50", 2.5, True) falls back for that key only, with a warning in the log. A provider that raises, or that returns something other than a mapping, falls back for both keys with a warning. Nothing is ever raised out of a resolver.

How parent and child agree on one budget

The old docstring explained why a request-scoped override was forbidden. The forked child re-resolves from os.environ, so an override only the parent could see would let the supervisor bound sit below the budget the child enforces. That would invert watchdog < native (+5) < outer (+10) < supervisor (+60).

This PR allows such an override and carries it across the fork explicitly, on the same terms as _pass_posture (ACE-101):

  1. tool_execute_sql resolves the organisation's limits once, before anything else, into execute_sql._statement_limits: one ContextVar holding both numbers, so a reader never sees one organisation's row cap beside another's deadline. The pin is reset in a finally.
  2. _resolve_row_cap / _resolve_timeout_s read the pin first, then the environment. The environment parsing moved unchanged into _row_cap_from_env / _timeout_s_from_env. Every bound calls the two resolvers, so all of them derive from the pinned budget: watchdog, native skew, outer bound (whose worker thread gets the pin through copy_context), and the fork path's supervisor.
  3. _pass_child_env writes AGAMI_SQL_MAX_ROWS / AGAMI_SQL_TIMEOUT_S into the child's environment, set to the resolved numbers. It writes them always, not only when an override applies, so the child parses a value the parent already resolved. The child has no provider and no pin, and reads exactly those keys back.

A test drives this across a real process boundary: the parent pins acme's limits, builds the child environment, and a real interpreter resolves the same row cap and timeout, below the parent's supervisor bound. Another test uses a provider whose answer changes on every call. It proves the provider is called once per call, and that the child environment, the supervisor bound and statement_limits() all see that one answer.

The docstrings that said the environment is the only source have been rewritten to describe this construction.

The tool description: built per caller at list-tools time

mcp_http's list_tools handler runs inside the request task, with _current_org_ctx already set, so the organisation is known when descriptions are served. The registry is still built once at import, with the import-time sentence kept as _EXECUTE_SQL_LIMITS_AT_IMPORT. At list-tools time, tool_description replaces exactly that sentence with the caller's numbers.

  • A consumer tool named execute_sql that does not contain our sentence passes through untouched, so the subtractive-only promise of the visibility hook still holds for consumer tools.
  • The handler now builds the list through run_blocking, because the provider is consumer code and usually a database read. That is ACE-048's reason for keeping blocking work off the event loop.
  • The stdio server (mcp_harness) calls the same function, so both servers state the same numbers.
  • A client keeps the tool list for its session, so a changed limit reaches new sessions. An existing session meets the new number in the resource_limit refusal, which already re-resolves the number on every call and now sees the pin. The admin screen should say this.

I chose this over stating the numbers on every response because it keeps what the client is told before it writes SQL (#326's point) correct for its organisation. test_execute_sql_states_limits.py and test_ace087_result_bound.py, including the assertion that max_rows never appears in the description, still pass. The sentence now opens "Limits in force for you" rather than "Limits on this deployment".

Decisions

  • No ceiling, but a timeout must be representable. A time limit is refused, falling back to the deployment value with the usual warning, when timeout_s + 60 (the supervisor's slack, the largest bound derived from it) is at or above threading.TIMEOUT_MAX (about 9.2 billion seconds). This is not a policy limit. At that size threading.Timer, Thread.join and the supervisor's wait raise OverflowError. The outer bound raised it after starting its worker and before counting the abandonment, so an administrator could have switched off the abandoned-worker cap. The rule is applied in tools._provider_limit and in execute_sql._timeout_s_from_env, through one helper, _timeout_is_representable. The row cap has no equivalent constraint.
  • Evaluation runs apply the organisation's limits. golden_run._run_item wraps each case's scoring in tools.pinned_statement_limits(org), so the answer key and the generated statement share one pin. tools is imported lazily. tools only reaches semantic_model lazily, so there is no cycle, but a module-level import would make loading the evaluation runner load the whole tool registry. If tools cannot be imported, no provider can be registered either, so an unpinned run is correct rather than degraded.
  • statement_limits(org_id=...) inside a call asks the provider again. The pin records numbers, not which organisation they belong to. Omit the argument to get what the current call enforces. This is documented and tested.
  • Tool listing without a provider stays on the event loop. Nothing blocks in that case, so the thread hop is skipped.

Availability note: accepted risk, no ceiling

With no ceiling, an administrator can set a very long statement time limit. One async worker serves every organisation on a shared server, so an organisation that sets, say, an hour and runs slow statements holds worker threads, and on an injected executor abandoned-worker slots, for that long, reducing capacity for everyone else. The owner has accepted this risk. test_there_is_no_ceiling pins the decision, so adding a clamp later is a deliberate change rather than a quiet one.

Tests

New: tests/test_per_org_statement_limits.py

  • the override reaches both in-process resolvers, and only inside the call
  • two organisations in sequence do not leak into each other; the pin is released when the call raises
  • tool_execute_sql holds the caller's limits for the whole call, and asks the provider exactly once
  • fork path: the child environment carries the organisation's numbers, and the supervisor bound is derived from them
  • a real forked interpreter resolves the parent's budget
  • fallbacks: no provider, None, {}, None values, invalid values (with a warning), a provider that raises, a non-mapping, a non-callable refused at registration
  • one key overridden alone; no ceiling
  • statement_limits() for the current and a named organisation; statement_limit_defaults()
  • description: the caller's numbers, no max_rows, other tools and consumer descriptions untouched, and the real HTTP tools/list handler serving different numbers to two organisations
  • the resource_limit refusal names the organisation's row cap
  • create_app registers Adapters.statement_limits, and building an app without one clears it (checked by the limits that then apply, not only by the registry)
  • the real child is checked to enforce exactly the number the parent's supervisor bound was derived from, not merely a smaller one
  • the outer bound and the watchdog fire on the pinned 1-second timeout while the deployment's is 30 seconds
  • a timeout at threading.TIMEOUT_MAX falls back, from the provider and from the environment, with a warning; the largest armable value is still accepted
  • the abandoned-worker cap still holds under an unrepresentable environment timeout: the first stuck call is abandoned and counted rather than raising OverflowError, and the next one is refused as saturated
  • listing tools without a provider does not hop threads
  • naming an organisation inside a call asks the provider again

Added to tests/test_golden_run.py: both statements of an evaluation case are scored under the named organisation's limits, the provider is asked for that organisation, and no pin is left behind.

Updated: tests/test_ace038_timeout.py

  • _statement_limits joins the structural "configuration surface" allowlist, with a comment giving the terms it is allowed on
  • the supervisor test now permits the two budget keys in the child environment, and asserts their values

plugins/agami/lib/execute_sql.py is byte-identical to the core copy.

Full suite

Run at the first commit (before the review fixes):

39 failed, 5643 passed, 12 skipped. The 39 are the known failures already on main: test_admin_model.py (26), test_admin_activity.py (6), test_model_deploy.py (2), test_prompt_examples_serving.py (2), test_ace110_query_basis.py, test_cloud_neutral_config.py and test_model_store_roundtrip.py (1 each). None of them mention statement limits, AGAMI_SQL_* or tool descriptions, and this change adds no new failure.

After the review fixes, the targeted suites (per-org limits, golden run and golden eval end-to-end, ace038 timeout, ace087 result bound, execute_sql limits, plugin copy parity) gave 1 failed, 1231 passed, 1 skipped. The one failure is test_ace110_query_basis.py::test_the_stored_basis_reaches_the_reader, which was already among the 39 failures on main.

Not covered

  • The CLI path (python -m execute_sql run by the local skill) reads the environment only, as before.

🤖 Generated with Claude Code

ashwin-agami and others added 2 commits September 14, 2026 15:32
…gine 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
@ashwin-agami
ashwin-agami marked this pull request as ready for review September 14, 2026 22:54
Copilot AI lite review requested due to automatic review settings September 14, 2026 22:54

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Unresolved moderate findings remain around timeout and row-cap representability, context preservation, provider fallback handling, and stdio listing coverage.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds per-organisation SQL row caps and statement timeouts through a provider seam, with limits pinned across execution, forks, MCP descriptions, and evaluations.

Changes:

  • Adds provider registration, validation, defaults, and request-scoped pinning.
  • Propagates limits through executors, child processes, and evaluations.
  • Personalizes HTTP/stdio tool descriptions and adds coverage/documentation.
File summaries
File Summary and final review notes
tests/test_per_org_statement_limits.py Adds comprehensive per-organisation limit coverage.
tests/test_golden_run.py Covers evaluation limits.
tests/test_ace038_timeout.py Covers fork and timeout budget behavior.
plugins/agami/lib/execute_sql.py Moderate (1 vote): handle native timeout representability. Nit (2 votes): correct copied-context documentation.
packages/agami-core/src/tools.py Moderate (1 vote): validate row-cap representability. Moderate (1 vote): guard custom mapping access. Nit (2 votes): neutralize deployment-cap wording. Nit (1 vote): clarify timeout representability in the provider contract.
packages/agami-core/src/semantic_model/golden_run.py Applies organisation-scoped limits during evaluations.
packages/agami-core/src/ports.py Adds the statement-limits adapter seam.
packages/agami-core/src/mcp_http.py Moderate (3 votes): preserve request/event-loop context for visibility predicates.
packages/agami-core/src/mcp_harness.py Moderate (1 vote): add end-to-end stdio listing coverage for overridden limits.
packages/agami-core/src/execute_sql.py Moderate (2 votes): handle native timeout representability. Nit (2 votes): correct documentation about worker context propagation.
CHANGELOG.md Documents the feature.
Review details

Suppressed comments (5)

packages/agami-core/src/mcp_harness.py:88

  • The new stdio listing path is not verified for per-caller limits: test_mcp_harness.py::test_initialize_and_tools_list only checks tool names, while the added tests cover tool_description directly and the low-level HTTP handler. Since this change explicitly promises that both transports advertise the same numbers, add an end-to-end stdio tools/list assertion under an overridden deployment limit.
            # 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"],

packages/agami-core/src/tools.py:1913

  • This public provider contract says either setting accepts any positive integer, but timeout_s is intentionally rejected when its supervisor bound cannot be represented by threading. That makes the documented API inaccurate for the very edge case this change adds; distinguish the absence of a policy ceiling from the platform-representability validation.
    ``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

packages/agami-core/src/tools.py:1934

  • This accepts every positive Python int as a usable row cap, but the execution paths pass cap + 1 to DB-API fetchmany (and the Postgres path also assigns it to cursor.itersize). Values above the driver's C integer range can therefore fail the statement with OverflowError instead of falling back or producing a row-limit refusal. Please add a representability check for the transfer APIs (distinct from a policy ceiling) and apply it consistently to provider/environment values.
    usable = isinstance(value, int) and not isinstance(value, bool) and value > 0

packages/agami-core/src/tools.py:1985

  • The mapping access is outside the provider-exception boundary. A provider is allowed to return any Mapping, but a lazy/custom mapping whose __bool__ or get() raises will escape _effective_statement_limits, contradicting the contract that provider failures fall back for both keys and never escape a resolver. Snapshot/read the mapping inside the same guarded fallback path.
    for key in _STATEMENT_LIMIT_KEYS:
        value = _provider_limit(org_id, key, (supplied or {}).get(key))

plugins/agami/lib/execute_sql.py:1651

  • This representability check only covers Python's timer/join/supervisor bounds, but the same timeout is also programmed into native engine limits. For example, _run_postgres multiplies it by 1000 for PostgreSQL statement_timeout; a value accepted here can overflow that millisecond setting and make SET LOCAL fail before the query runs, producing a generic executor failure instead of a bounded statement. Add native-bound representability handling (or skip an unrepresentable native backstop while retaining the watchdog) in both byte-identical copies.
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
  • Files reviewed: 11/11 changed files
  • Comments generated: 5
  • Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +1641 to +1651
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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #342: timeouts over 604,800 s (seven days, Snowflake's native maximum and the smallest among the engines, so also safe for Postgres's int32 ms setting) are treated as unusable and fall back to the deployment value.

Comment on lines +508 to +510
if not has_statement_limits_provider():
return _listed()
return await run_blocking(_listed)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #342: the visibility predicate now runs in the request task; only the description/provider work is offloaded. A test pins it.

Comment on lines +2706 to +2708
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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #342: comment corrected. The copied context is load-bearing, because the engine functions in the worker call both resolvers.

Comment on lines 3586 to +3590
"do an out-of-scope table or column, a per-statement deadline, and a result larger "
"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`.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #342: the wording is now 'the row limit', and the following sentence states the caller's number.

Comment on lines +2706 to +2708
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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #342: the mirror is byte-identical to the corrected core copy.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@ashwin-agami
ashwin-agami enabled auto-merge (squash) September 14, 2026 23:10
@ashwin-agami
ashwin-agami merged commit f8823f4 into main Sep 14, 2026
9 checks passed
@ashwin-agami
ashwin-agami deleted the per-org-statement-limits branch September 14, 2026 23:14
@github-actions github-actions Bot locked and limited conversation to collaborators Sep 14, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants