diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 7d5bea20..88fc7cf9 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -736,10 +736,16 @@ def _model_safety(sql: str, profile: str, area: str | None): sys.stderr.write(f"[agami] could not load semantic model; skipping safety pass: {e}\n") return sql, None + # Build the shared guard context ONCE — parse the SQL + build each model index a single + # time — and thread it through the battery below, instead of every guard re-parsing and + # rebuilding its index (audit P2 / ACE-045). Behaviour-preserving: a guard given `ctx` + # returns the same verdict as one that builds its own. + ctx = RT.build_guard_context(sql, org) + # Table-scope guard — a query may only reference tables the semantic model # declares; any other table in the connected database is refused. Runs FIRST # so the fan/chasm and sensitive checks below only evaluate in-scope tables. - ts = RT.check_table_scope(sql, org) + ts = RT.check_table_scope(sql, org, ctx=ctx) if ts.action == "refuse": json.dump({"error": {"kind": "table_out_of_scope", "tables": ts.offending_tables, "reason": ts.reason, "suggestion": ts.suggestion}}, sys.stderr) @@ -748,7 +754,7 @@ def _model_safety(sql: str, profile: str, area: str | None): # SELECT * ban — force every projected column to be named, so the column-scope # guard below can check what is actually returned (and nothing hides behind *). - star = RT.check_no_select_star(sql) + star = RT.check_no_select_star(sql, ctx=ctx) if star.action == "refuse": json.dump({"error": {"kind": "select_star", "reason": star.reason, "suggestion": star.suggestion}}, sys.stderr) @@ -757,14 +763,14 @@ def _model_safety(sql: str, profile: str, area: str | None): # Column-scope guard — a column that binds to a declared table must be one that # table declares (a hallucinated column, or a physical column the model excluded). - cs = RT.check_column_scope(sql, org) + cs = RT.check_column_scope(sql, org, ctx=ctx) if cs.action == "refuse": json.dump({"error": {"kind": "column_out_of_scope", "columns": cs.columns, "reason": cs.reason, "suggestion": cs.suggestion}}, sys.stderr) sys.stderr.write("\n") return sql, 1 - pf = RT.pre_flight_check(sql, org) + pf = RT.pre_flight_check(sql, org, ctx=ctx) if pf.risk and pf.action == "refuse": json.dump({"error": {"kind": "preflight_refused", "risk": pf.risk, "reason": pf.reason, "suggestion": pf.suggestion, @@ -774,19 +780,20 @@ def _model_safety(sql: str, profile: str, area: str | None): if pf.risk and pf.action == "auto_rewrite" and pf.rewritten_sql: sys.stderr.write(f"[agami] auto-corrected {pf.risk}: ran rewritten SQL. {pf.reason}\n") sql = pf.rewritten_sql + ctx = RT.build_guard_context(sql, org) # SQL changed -> refresh the shared context # Sensitive-column (PII) guard — refuse to PROJECT raw sensitive values. Same # deterministic chokepoint as the fan/chasm pre-flight, so the agami-query skill, # the local MCP server, and cron all protect PII identically (not just whichever # path happened to read a prose rule). Aggregates / filters / joins are allowed. - sens = RT.check_sensitive_projection(sql, org) + sens = RT.check_sensitive_projection(sql, org, ctx=ctx) if sens.action == "refuse": json.dump({"error": {"kind": "sensitive_columns", "columns": sens.columns, "reason": sens.reason, "suggestion": sens.suggestion}}, sys.stderr) sys.stderr.write("\n") return sql, 1 - new_sql, applied = RT.apply_default_filters(sql, org, area=area) + new_sql, applied = RT.apply_default_filters(sql, org, area=area, ctx=ctx) if applied: sys.stderr.write(f"[agami] applied default_filters: {applied}\n") sql = new_sql diff --git a/packages/agami-core/src/mcp_http.py b/packages/agami-core/src/mcp_http.py index 00885ed0..48bc0db5 100644 --- a/packages/agami-core/src/mcp_http.py +++ b/packages/agami-core/src/mcp_http.py @@ -46,6 +46,7 @@ SERVER_INSTRUCTIONS, SERVER_NAME, TOOLS, + _current_org_ctx, bootstrap_paths, record_tool_call, server_version, @@ -73,6 +74,15 @@ def _actor_from_scope(scope: dict, auth: AuthProvider) -> str | None: return None +def _org_id_from_scope(scope: dict) -> str | None: + """The resolved org id for this /mcp request — the org the auth middleware attached to the ASGI scope + state (`request.state.org`). None under presence auth / single-tenant, where the tool layer falls back + to AGAMI_ORG_ID / 'local'. Keeps the per-process model cache tenant-safe (ACE-045): cache entries key on + this, so under a multi-tenant resolver one org never gets another's cached model.""" + org = (scope.get("state") or {}).get("org") + return getattr(org, "id", None) + + # The brand assets (logo, provider icons, favicon) served at /static — packaged alongside this module. _STATIC_DIR = Path(__file__).resolve().parent / "static" @@ -362,15 +372,17 @@ def create_app(extra_tools: dict | None = None, adapters: Adapters | None = None ) async def handle_mcp(scope, receive, send): - # Set the actor for this request's tool calls, then run the MCP dispatch in the same task so the - # contextvar reaches `_call_tool`. Prefer the principal the middleware validated (on scope state); - # fall back to re-validating the bearer from the scope headers if it didn't propagate. + # Set the actor + resolved org for this request's tool calls, then run the MCP dispatch in the same + # task so the contextvars reach `_call_tool` and the per-process model cache. Prefer what the auth + # middleware attached to the scope state; the actor falls back to re-validating the bearer. actor = _actor_from_scope(scope, auth_provider) token = _actor_ctx.set(actor) + org_token = _current_org_ctx.set(_org_id_from_scope(scope)) try: await session_manager.handle_request(scope, receive, send) finally: _actor_ctx.reset(token) + _current_org_ctx.reset(org_token) @contextlib.asynccontextmanager async def lifespan(_app: Starlette): diff --git a/packages/agami-core/src/semantic_model/runtime.py b/packages/agami-core/src/semantic_model/runtime.py index 9ff9a78f..a9cc0733 100644 --- a/packages/agami-core/src/semantic_model/runtime.py +++ b/packages/agami-core/src/semantic_model/runtime.py @@ -58,6 +58,57 @@ # exists in .. Injected so runtime stays DB-agnostic. Prober = Callable[[str, str, str], bool] + +# --------------------------------------------------------------------------- +# Per-invocation guard context (ACE-045) +# +# The _model_safety battery (execute_sql.py) runs ~6 guards that EACH re-parse the SQL +# (sqlglot ×6) and rebuild their model index from scratch. `GuardContext` does that +# shared work ONCE — the SQL parsed once, each index built once — and is threaded through +# the guards via an optional `ctx=`. A guard given `ctx` returns the SAME verdict as one +# that builds its own (behaviour-preserving); `ctx=None` keeps the standalone callers +# (e.g. cli.py) working unchanged. `tree` is None when the SQL doesn't parse — guards then +# degrade to allow, exactly as the inline parse-and-except did before. +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class GuardContext: + sql: str + tree: "exp.Expression | None" + column_index: "dict[str, dict[str, Column]]" + cardinality_index: "list[Relationship]" + sensitive_by_table: "tuple[dict[str, set[str]], set[str]]" + model_table_index: "dict[str, tuple]" + + +def _parse_sql(sql: str) -> "exp.Expression | None": + """Parse SQL for the guard battery; None if sqlglot is unavailable or the SQL does not + parse (guards degrade to allow). Centralized so a GuardContext parses exactly once.""" + if not _HAVE_SQLGLOT: + return None + try: + return sqlglot.parse_one(sql, error_level="ignore") + except Exception: + return None + + +def build_guard_context(sql: str, org: Organization) -> "GuardContext | None": + """Parse `sql` once and build each guard index once, so the _model_safety battery shares + them instead of every guard redoing the work (audit P2 / ACE-045). Returns None when sqlglot + is unavailable: every guard then short-circuits to allow before it touches the context, so + building the indices would be pure wasted work in that fallback path.""" + if not _HAVE_SQLGLOT: + return None + return GuardContext( + sql=sql, + tree=_parse_sql(sql), + column_index=_column_index(org), + cardinality_index=_cardinality_index(org), + sensitive_by_table=_sensitive_by_table(org), + model_table_index=_model_table_index(org), + ) + # Ambiguity threshold — "ask, don't guess" when top-two are within this delta. AMBIGUITY_DELTA = 0.15 @@ -447,19 +498,26 @@ def _output_selects(node: "exp.Expression") -> list["exp.Select"]: return [] -def check_sensitive_projection(sql: str, org: Organization) -> SensitiveCheckResult: +def check_sensitive_projection(sql: str, org: Organization, + ctx: "GuardContext | None" = None) -> SensitiveCheckResult: """Refuse a query that PROJECTS a `sensitive` column's raw values; allow the column in COUNT, filters, GROUP BY, and joins. Degrades to allow when sqlglot - is unavailable or the SQL doesn't parse (same posture as the fan/chasm pass).""" + is unavailable or the SQL doesn't parse (same posture as the fan/chasm pass). + + `ctx` (ACE-045): reuse the once-parsed tree + once-built sensitive index instead of + redoing both; `ctx=None` keeps the standalone path byte-identical.""" if not _HAVE_SQLGLOT: return SensitiveCheckResult("allow") - by_table, allnames = _sensitive_by_table(org) + by_table, allnames = ctx.sensitive_by_table if ctx is not None else _sensitive_by_table(org) if not allnames: return SensitiveCheckResult("allow") - try: - tree = sqlglot.parse_one(sql, error_level="ignore") - except Exception: - return SensitiveCheckResult("allow") + if ctx is not None: + tree = ctx.tree + else: + try: + tree = sqlglot.parse_one(sql, error_level="ignore") + except Exception: + return SensitiveCheckResult("allow") # A set operation (UNION/INTERSECT/EXCEPT) parses to exp.SetOperation, not # exp.Select — gate on "contains a SELECT" and scan every OUTPUT-bearing arm, else # `… UNION SELECT ssn FROM customers` would project a sensitive column past this gate. @@ -528,7 +586,8 @@ def as_dict(self) -> dict[str, Any]: "reason": self.reason, "suggestion": self.suggestion} -def check_table_scope(sql: str, org: Organization) -> TableScopeResult: +def check_table_scope(sql: str, org: Organization, + ctx: "GuardContext | None" = None) -> TableScopeResult: """Refuse a query that references a table not declared in the semantic model. Only *physical* table references count: CTE names (defined by WITH) and @@ -545,13 +604,16 @@ def check_table_scope(sql: str, org: Organization) -> TableScopeResult: """ if not _HAVE_SQLGLOT: return TableScopeResult("allow") - allow = {name.lower() for name in _model_table_index(org)} + allow = {name.lower() for name in (ctx.model_table_index if ctx is not None else _model_table_index(org))} if not allow: return TableScopeResult("allow") - try: - tree = sqlglot.parse_one(sql, error_level="ignore") - except Exception: - return TableScopeResult("allow") + if ctx is not None: + tree = ctx.tree + else: + try: + tree = sqlglot.parse_one(sql, error_level="ignore") + except Exception: + return TableScopeResult("allow") # A set operation (UNION/INTERSECT/EXCEPT) parses to exp.Union, not exp.Select, # so gate on "contains a SELECT" rather than "is a SELECT" — otherwise every # set-operation arm would bypass the guard. A non-SELECT statement has no SELECT @@ -603,7 +665,7 @@ def as_dict(self) -> dict[str, Any]: return {"action": self.action, "reason": self.reason, "suggestion": self.suggestion} -def check_no_select_star(sql: str) -> StarCheckResult: +def check_no_select_star(sql: str, ctx: "GuardContext | None" = None) -> StarCheckResult: """Refuse a query whose projection list contains `*` or `t.*`. A star defeats column-level scoping (an undeclared column hides behind it) and @@ -618,10 +680,13 @@ def check_no_select_star(sql: str) -> StarCheckResult: """ if not _HAVE_SQLGLOT: return StarCheckResult("allow") - try: - tree = sqlglot.parse_one(sql, error_level="ignore") - except Exception: - return StarCheckResult("allow") + if ctx is not None: + tree = ctx.tree + else: + try: + tree = sqlglot.parse_one(sql, error_level="ignore") + except Exception: + return StarCheckResult("allow") if tree is None or tree.find(exp.Select) is None: return StarCheckResult("allow") for select in tree.find_all(exp.Select): @@ -648,7 +713,8 @@ def as_dict(self) -> dict[str, Any]: "reason": self.reason, "suggestion": self.suggestion} -def check_column_scope(sql: str, org: Organization) -> ColumnScopeResult: +def check_column_scope(sql: str, org: Organization, + ctx: "GuardContext | None" = None) -> ColumnScopeResult: """Refuse a query that references a column not declared on the table it binds to. Strict where a column visibly binds to a declared physical table — qualified by @@ -667,13 +733,16 @@ def check_column_scope(sql: str, org: Organization) -> ColumnScopeResult: """ if not _HAVE_SQLGLOT: return ColumnScopeResult("allow") - colidx = _column_index(org) + colidx = ctx.column_index if ctx is not None else _column_index(org) if not colidx: return ColumnScopeResult("allow") - try: - tree = sqlglot.parse_one(sql, error_level="ignore") - except Exception: - return ColumnScopeResult("allow") + if ctx is not None: + tree = ctx.tree + else: + try: + tree = sqlglot.parse_one(sql, error_level="ignore") + except Exception: + return ColumnScopeResult("allow") if tree is None or tree.find(exp.Select) is None: return ColumnScopeResult("allow") @@ -812,7 +881,8 @@ def _many_side_facing_one(rels: list[Relationship], table: str, dim: str) -> boo return False -def pre_flight_check(sql: str, org: Organization) -> PreFlightResult: +def pre_flight_check(sql: str, org: Organization, + ctx: "GuardContext | None" = None) -> PreFlightResult: """Detect fan-trap / chasm-trap and decide rewrite-vs-refuse-vs-allow. A set operation (UNION/INTERSECT/EXCEPT) parses to exp.SetOperation, not exp.Select; @@ -822,17 +892,17 @@ def pre_flight_check(sql: str, org: Organization) -> PreFlightResult: unavailable, the SQL doesn't parse, or it contains no SELECT.""" if not _HAVE_SQLGLOT: return PreFlightResult(None, "allow", sql, reason="sqlglot unavailable; skipped") - try: - tree = sqlglot.parse_one(sql, error_level="ignore") - except Exception as e: - return PreFlightResult(None, "allow", sql, reason=f"unparseable; skipped ({e})") + # Parse via the same centralized helper the ctx path used (ACE-045), so a ctx and a non-ctx + # call are byte-identical: _parse_sql swallows an unparseable statement to None exactly as a + # prebuilt ctx.tree would be None, and both then report the one "no SELECT; skipped" reason. + tree = ctx.tree if ctx is not None else _parse_sql(sql) if tree is None or tree.find(exp.Select) is None: return PreFlightResult(None, "allow", sql, reason="no SELECT; skipped") if isinstance(tree, exp.Select): - return _preflight_select(tree, org, sql, allow_rewrite=True) + return _preflight_select(tree, org, sql, allow_rewrite=True, ctx=ctx) # Set operation: analyze each arm; a trap in any arm inflates that arm's aggregate. for arm in _output_selects(tree): - res = _preflight_select(arm, org, arm.sql(), allow_rewrite=False) + res = _preflight_select(arm, org, arm.sql(), allow_rewrite=False, ctx=ctx) if res.risk and res.action == "refuse": # tie the arm's diagnosis back to the full set-operation query return PreFlightResult(res.risk, "refuse", sql, reason=res.reason, @@ -840,11 +910,15 @@ def pre_flight_check(sql: str, org: Organization) -> PreFlightResult: return PreFlightResult(None, "allow", sql, reason="no fan/chasm or aggregation issue in any arm") -def _preflight_select(tree: "exp.Select", org: Organization, sql: str, allow_rewrite: bool) -> PreFlightResult: +def _preflight_select(tree: "exp.Select", org: Organization, sql: str, allow_rewrite: bool, + ctx: "GuardContext | None" = None) -> PreFlightResult: """Fan/chasm + aggregation-semantics analysis of a SINGLE SELECT. `sql` is that select's own text (used for the join rewrite + messages). When `allow_rewrite` is - False (a set-operation arm), a rewriteable fan trap is refused, not rewritten.""" - rels = _cardinality_index(org) + False (a set-operation arm), a rewriteable fan trap is refused, not rewritten. + + `ctx` supplies the shared cardinality/column indices (ACE-045); `tree` is always the + caller's own SELECT (a set-op arm ≠ `ctx.tree`), so only the indices come from `ctx`.""" + rels = ctx.cardinality_index if ctx is not None else _cardinality_index(org) tables_in_scope = _tables_in_scope(tree) # alias -> table table_set = set(tables_in_scope.values()) @@ -934,7 +1008,7 @@ def _preflight_select(tree: "exp.Select", org: Organization, sql: str, allow_rew # No structural (join) trap. Now the SEMANTIC checks the fan/chasm detector is # blind to (scorecard #4): aggregation-class violations (#2) and semi-additive # rollups over time (#3) — these need NO join, so cardinality analysis can't see them. - semantic = _check_aggregation_semantics(tree, org, tables_in_scope, sql) + semantic = _check_aggregation_semantics(tree, org, tables_in_scope, sql, ctx=ctx) if semantic is not None: return semantic @@ -1028,9 +1102,10 @@ def _groups_by_time(tree: "exp.Select", scope: dict[str, str], def _check_aggregation_semantics( - tree: "exp.Select", org: Organization, scope: dict[str, str], sql: str + tree: "exp.Select", org: Organization, scope: dict[str, str], sql: str, + ctx: "GuardContext | None" = None, ) -> Optional[PreFlightResult]: - colidx = _column_index(org) + colidx = ctx.column_index if ctx is not None else _column_index(org) # --- #2: aggregation-class violations (SUM of a rate/id, AVG of an id) --- for select_expr in tree.expressions: @@ -1428,6 +1503,7 @@ def apply_default_filters( *, area: Optional[str] = None, params: Optional[dict[str, str]] = None, + ctx: "GuardContext | None" = None, ) -> tuple[str, list[str]]: """Conservatively AND each in-scope table's default_filters into the SQL's WHERE. @@ -1443,10 +1519,16 @@ def apply_default_filters( params = params or {} if not _HAVE_SQLGLOT: return sql, [] - try: - tree = sqlglot.parse_one(sql, error_level="ignore") - except Exception: - return sql, [] + if ctx is not None: + # apply_default_filters MUTATES its tree to inject WHEREs; work on a COPY so the + # shared ctx.tree the read-only guards used stays pristine (ACE-045). ctx must be + # built from this same `sql`, which _model_safety guarantees. + tree = ctx.tree.copy() if ctx.tree is not None else None + else: + try: + tree = sqlglot.parse_one(sql, error_level="ignore") + except Exception: + return sql, [] if not isinstance(tree, exp.Select): return sql, [] diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index e10b6f34..eefabef8 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -29,6 +29,7 @@ import sys import time from collections.abc import Callable +from contextvars import ContextVar from pathlib import Path from typing import Any @@ -240,7 +241,7 @@ def _resolve_units(profile: str, sql: str) -> dict[str, str]: model deps (pydantic/sqlglot) aren't installed — execute_sql stays pure-stdlib; numbers still format exactly via units.py, just without a currency symbol.""" try: - org = _load_org(profile) + org = get_cached_org(profile) from semantic_model import runtime as RT return RT.resolve_result_units(org, sql) @@ -272,6 +273,47 @@ def _model_version(profile: str) -> str | None: return None +# The org id for the current request's tool calls (ACE-045). The HTTP server sets this per request from +# the OrgResolver-resolved org; unset (stdio / single-tenant) it falls back to AGAMI_ORG_ID / "local". +_current_org_ctx: ContextVar[str | None] = ContextVar("agami_current_org_id", default=None) + + +def _current_org_id() -> str: + """The org id to scope this process's model cache by: the request's resolved org when the HTTP server + set it (per-request under a multi-tenant resolver), else AGAMI_ORG_ID / 'local' (single-tenant / stdio).""" + return _current_org_ctx.get() or os.environ.get("AGAMI_ORG_ID") or "local" + + +# Per-process semantic-model cache (ACE-045). The long-lived server loads the whole model 2-3x per query +# (_resolve_units + _resolve_receipt) and re-loads it every query; caching serves it warm across queries and +# users. Keyed (org_id, datasource, model_version): org-scoped so a multi-tenant server never serves one org's +# model to another, and invalidated when the model version bumps. The execute_sql subprocess is a fresh process +# per query and does NOT share this (its win is the Slice-1 GuardContext, not a cross-query cache). +_ORG_CACHE: dict[tuple[str, str, "str | None"], Any] = {} + + +def get_cached_org(profile: str): + """Load the semantic model for `profile`, cached per process and keyed (org, datasource, version). + Reuses one Organization across the loads within a query AND across queries, until the model version + changes; a cache miss falls back to a fresh `_load_org`.""" + version = _model_version(profile) # cheap: one DB row / dir listing, not a full model load + if version is None: + # No version to detect a model change by (e.g. file mode with no snapshot) — don't cache, + # so we can never serve a stale model. The DB-backed server always has a version. + return _load_org(profile) + org_id = _current_org_id() + key = (org_id, profile, version) + cached = _ORG_CACHE.get(key) + if cached is not None: + return cached + org = _load_org(profile) + # Drop any stale (org, datasource) entry at a previous version so the cache stays bounded. + for stale in [k for k in _ORG_CACHE if k[0] == org_id and k[1] == profile and k != key]: + del _ORG_CACHE[stale] + _ORG_CACHE[key] = org + return org + + def _domain_memory(profile: str) -> tuple[str, str | None]: """(ORGANIZATION.md text, USER_MEMORY.md text) for the domain-context block — from the DB when AGAMI_DB_URL is set (no file read at runtime), else from disk.""" @@ -298,7 +340,7 @@ def _resolve_receipt(profile: str, sql: str) -> dict | None: touch / what's unapproved' panel is identical in Claude Code and Claude Desktop. Returns None only if the model deps aren't importable (execute_sql stays usable).""" try: - org = _load_org(profile) + org = get_cached_org(profile) from semantic_model import runtime as RT return RT.assemble_receipt(org, sql, model_version=_model_version(profile)) @@ -647,7 +689,7 @@ def tool_get_datasource_schema(args: dict[str, Any]) -> str: """ profile = resolve_profile(args.get("datasource")) try: - org = _load_org(profile) + org = get_cached_org(profile) except FileNotFoundError as e: return json.dumps({"error": {"kind": "not_found", "remediation": str(e)}}, indent=2) except ImportError: diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 7d5bea20..88fc7cf9 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -736,10 +736,16 @@ def _model_safety(sql: str, profile: str, area: str | None): sys.stderr.write(f"[agami] could not load semantic model; skipping safety pass: {e}\n") return sql, None + # Build the shared guard context ONCE — parse the SQL + build each model index a single + # time — and thread it through the battery below, instead of every guard re-parsing and + # rebuilding its index (audit P2 / ACE-045). Behaviour-preserving: a guard given `ctx` + # returns the same verdict as one that builds its own. + ctx = RT.build_guard_context(sql, org) + # Table-scope guard — a query may only reference tables the semantic model # declares; any other table in the connected database is refused. Runs FIRST # so the fan/chasm and sensitive checks below only evaluate in-scope tables. - ts = RT.check_table_scope(sql, org) + ts = RT.check_table_scope(sql, org, ctx=ctx) if ts.action == "refuse": json.dump({"error": {"kind": "table_out_of_scope", "tables": ts.offending_tables, "reason": ts.reason, "suggestion": ts.suggestion}}, sys.stderr) @@ -748,7 +754,7 @@ def _model_safety(sql: str, profile: str, area: str | None): # SELECT * ban — force every projected column to be named, so the column-scope # guard below can check what is actually returned (and nothing hides behind *). - star = RT.check_no_select_star(sql) + star = RT.check_no_select_star(sql, ctx=ctx) if star.action == "refuse": json.dump({"error": {"kind": "select_star", "reason": star.reason, "suggestion": star.suggestion}}, sys.stderr) @@ -757,14 +763,14 @@ def _model_safety(sql: str, profile: str, area: str | None): # Column-scope guard — a column that binds to a declared table must be one that # table declares (a hallucinated column, or a physical column the model excluded). - cs = RT.check_column_scope(sql, org) + cs = RT.check_column_scope(sql, org, ctx=ctx) if cs.action == "refuse": json.dump({"error": {"kind": "column_out_of_scope", "columns": cs.columns, "reason": cs.reason, "suggestion": cs.suggestion}}, sys.stderr) sys.stderr.write("\n") return sql, 1 - pf = RT.pre_flight_check(sql, org) + pf = RT.pre_flight_check(sql, org, ctx=ctx) if pf.risk and pf.action == "refuse": json.dump({"error": {"kind": "preflight_refused", "risk": pf.risk, "reason": pf.reason, "suggestion": pf.suggestion, @@ -774,19 +780,20 @@ def _model_safety(sql: str, profile: str, area: str | None): if pf.risk and pf.action == "auto_rewrite" and pf.rewritten_sql: sys.stderr.write(f"[agami] auto-corrected {pf.risk}: ran rewritten SQL. {pf.reason}\n") sql = pf.rewritten_sql + ctx = RT.build_guard_context(sql, org) # SQL changed -> refresh the shared context # Sensitive-column (PII) guard — refuse to PROJECT raw sensitive values. Same # deterministic chokepoint as the fan/chasm pre-flight, so the agami-query skill, # the local MCP server, and cron all protect PII identically (not just whichever # path happened to read a prose rule). Aggregates / filters / joins are allowed. - sens = RT.check_sensitive_projection(sql, org) + sens = RT.check_sensitive_projection(sql, org, ctx=ctx) if sens.action == "refuse": json.dump({"error": {"kind": "sensitive_columns", "columns": sens.columns, "reason": sens.reason, "suggestion": sens.suggestion}}, sys.stderr) sys.stderr.write("\n") return sql, 1 - new_sql, applied = RT.apply_default_filters(sql, org, area=area) + new_sql, applied = RT.apply_default_filters(sql, org, area=area, ctx=ctx) if applied: sys.stderr.write(f"[agami] applied default_filters: {applied}\n") sql = new_sql diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..07d9015a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,27 @@ +"""Shared pytest fixtures for agami-core tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "plugins" / "agami" / "scripts")) + + +@pytest.fixture(autouse=True) +def _reset_org_cache(): + """The per-process semantic-model cache (ACE-045) is module-global state; isolate every test from it + (and from a leaked current-org) so one test's cached model never bleeds into the next.""" + try: + import tools + except Exception: + yield + return + tools._ORG_CACHE.clear() + tools._current_org_ctx.set(None) + yield + tools._ORG_CACHE.clear() + tools._current_org_ctx.set(None) diff --git a/tests/test_guard_context.py b/tests/test_guard_context.py new file mode 100644 index 00000000..e3b0434b --- /dev/null +++ b/tests/test_guard_context.py @@ -0,0 +1,119 @@ +"""ACE-045 Slice 1: the shared `GuardContext` parses the SQL once and builds each model +index once, and every guard returns an identical verdict with or without a `ctx` — so +threading it through `_model_safety` is behaviour-preserving, only cheaper.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("pydantic") +pytest.importorskip("sqlglot") + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "plugins" / "agami" / "scripts")) + +from semantic_model import models as m # noqa: E402 +from semantic_model import runtime as rt # noqa: E402 + + +def _org(): + customers = m.Table( + name="customers", schema="public", storage_connection="c", grain=["id"], + default_filters=["{alias}.active = true"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="name", type="string"), + m.Column(name="email", type="string", sensitive=True), + ], + ) + orders = m.Table( + name="orders", schema="public", storage_connection="c", grain=["id"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="customer_id", type="integer"), + m.Column(name="total_amount", type="decimal", aggregation="additive"), + ], + ) + rels = [m.Relationship(from_table="orders", to_table="customers", + from_column="customer_id", to_column="id", + relationship="many_to_one")] + return m.Organization( + organization="Shop", + subject_areas=[m.SubjectArea(name="sales", tables_defined=[customers, orders], + relationships=rels)], + ) + + +_CLEAN = ("SELECT customers.name, COUNT(orders.id) AS n FROM customers " + "JOIN orders ON orders.customer_id = customers.id GROUP BY customers.name") + + +def test_build_context_parses_and_indexes_each_once(monkeypatch): + """build_guard_context does the shared work once; the guards given `ctx` add none.""" + org = _org() + counts = {"parse": 0, "_column_index": 0, "_cardinality_index": 0, + "_sensitive_by_table": 0, "_model_table_index": 0} + + real_parse = rt.sqlglot.parse_one + monkeypatch.setattr(rt.sqlglot, "parse_one", + lambda *a, **k: (counts.__setitem__("parse", counts["parse"] + 1) + or real_parse(*a, **k))) + for name in ("_column_index", "_cardinality_index", "_sensitive_by_table", "_model_table_index"): + real = getattr(rt, name) + + def wrapper(org, _real=real, _name=name): + counts[_name] += 1 + return _real(org) + + monkeypatch.setattr(rt, name, wrapper) + + # A query on `orders` only — which declares no default_filters — so apply_default_filters + # injects nothing and we isolate the "guards don't re-parse the query SQL" claim (a table + # WITH default_filters legitimately parses each filter fragment to inject it). + sql = "SELECT COUNT(orders.id) AS n FROM orders" + ctx = rt.build_guard_context(sql, org) + # Full battery WITH ctx — none of these should parse or rebuild an index again. + rt.check_table_scope(sql, org, ctx=ctx) + rt.check_no_select_star(sql, ctx=ctx) + rt.check_column_scope(sql, org, ctx=ctx) + rt.pre_flight_check(sql, org, ctx=ctx) + rt.check_sensitive_projection(sql, org, ctx=ctx) + rt.apply_default_filters(sql, org, ctx=ctx) + + assert counts == {"parse": 1, "_column_index": 1, "_cardinality_index": 1, + "_sensitive_by_table": 1, "_model_table_index": 1} + + +@pytest.mark.parametrize("sql", [ + _CLEAN, # allow + "SELECT * FROM orders", # star ban + "SELECT customers.email FROM customers", # sensitive projection refuse + "SELECT customers.bogus_col FROM customers", # column-scope refuse + "SELECT ghost.x FROM ghost", # table-scope refuse + "SELECT customers.name FROM customers", # allow + default_filter applied +]) +def test_verdict_parity_with_and_without_ctx(sql): + """Every guard returns byte-identical results whether it builds its own work or is + handed a shared ctx — the behaviour-preserving guarantee.""" + org = _org() + ctx = rt.build_guard_context(sql, org) + assert rt.check_table_scope(sql, org).as_dict() == rt.check_table_scope(sql, org, ctx=ctx).as_dict() + assert rt.check_no_select_star(sql).as_dict() == rt.check_no_select_star(sql, ctx=ctx).as_dict() + assert rt.check_column_scope(sql, org).as_dict() == rt.check_column_scope(sql, org, ctx=ctx).as_dict() + assert rt.pre_flight_check(sql, org).as_dict() == rt.pre_flight_check(sql, org, ctx=ctx).as_dict() + assert (rt.check_sensitive_projection(sql, org).as_dict() + == rt.check_sensitive_projection(sql, org, ctx=ctx).as_dict()) + assert rt.apply_default_filters(sql, org) == rt.apply_default_filters(sql, org, ctx=ctx) + + +def test_unparseable_sql_ctx_tree_is_none_and_guards_allow(): + """A GuardContext over unparseable SQL carries tree=None; guards degrade to allow, + matching the standalone path.""" + org = _org() + bad = "NOT SQL AT ALL ;;;" + ctx = rt.build_guard_context(bad, org) + assert rt.check_table_scope(bad, org, ctx=ctx).action == "allow" + assert rt.check_no_select_star(bad, ctx=ctx).action == "allow" diff --git a/tests/test_org_cache.py b/tests/test_org_cache.py new file mode 100644 index 00000000..152e415e --- /dev/null +++ b/tests/test_org_cache.py @@ -0,0 +1,76 @@ +"""ACE-045 Slice 2: get_cached_org serves the semantic model warm across queries (one load, +then cache hits), reloads when the model version bumps, and is ORG-SCOPED — a multi-tenant +server never serves one org's model to another, even on the same datasource name.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "plugins" / "agami" / "scripts")) + +import tools # noqa: E402 + +# Cache isolation between tests is handled once in tests/conftest.py::_reset_org_cache (autouse) — +# no per-file duplicate here. + + +def test_loads_once_then_serves_warm(monkeypatch): + calls = {"load": 0} + monkeypatch.setattr(tools, "_model_version", lambda profile: "v1") + + def fake_load(profile): + calls["load"] += 1 + return {"org": profile} + + monkeypatch.setattr(tools, "_load_org", fake_load) + + first = tools.get_cached_org("sales") + second = tools.get_cached_org("sales") + assert calls["load"] == 1 # query #2 served warm from the cache + assert first is second # the same cached Organization object + + +def test_reloads_after_version_bump(monkeypatch): + calls = {"load": 0} + versions = iter(["v1", "v1", "v2"]) + monkeypatch.setattr(tools, "_model_version", lambda profile: next(versions)) + + def fake_load(profile): + calls["load"] += 1 + return object() + + monkeypatch.setattr(tools, "_load_org", fake_load) + + tools.get_cached_org("sales") # v1 -> load + tools.get_cached_org("sales") # v1 -> warm + tools.get_cached_org("sales") # v2 -> reload + assert calls["load"] == 2 + assert len(tools._ORG_CACHE) == 1 # stale v1 entry evicted, only v2 remains + + +def test_org_scoped_no_cross_tenant(monkeypatch): + monkeypatch.setattr(tools, "_model_version", lambda profile: "v1") + + def fake_load(profile): + return {"org_id": tools._current_org_id(), "profile": profile} + + monkeypatch.setattr(tools, "_load_org", fake_load) + + tools._current_org_ctx.set("orgA") + a = tools.get_cached_org("sales") + tools._current_org_ctx.set("orgB") + b = tools.get_cached_org("sales") # SAME datasource name, different org + + assert a is not b + assert a["org_id"] == "orgA" and b["org_id"] == "orgB" + # org A never receives org B's cached model + tools._current_org_ctx.set("orgA") + assert tools.get_cached_org("sales") is a + + +def test_default_org_id_falls_back_to_local(monkeypatch): + monkeypatch.delenv("AGAMI_ORG_ID", raising=False) + tools._current_org_ctx.set(None) + assert tools._current_org_id() == "local"