ACE-045: per-process model cache + single SQL parse + shared guard indices

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

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):

packages/agami-core/src/semantic_model/runtime.py

@@ -58,6 +58,57 @@
 # exists in <table>.<column>. 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,29 +892,33 @@ 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,
                                    suggestion=res.suggestion, triggering_joins=res.triggering_joins)
     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, []