From f20114c4f977c1a3f6cdbb4a57a16923a341c726 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:02:13 -0700 Subject: [PATCH 01/11] ACE-041 slice1: failing guard tests for maskable/index classification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TEST-FIRST checkpoint. Adds tests/test_sensitive_projection_maskable.py — a broad matrix asserting that check_sensitive_projection classifies each offending sensitive projection as maskable (with its 0-based output-column index) vs must-refuse (buried in expression/function/scalar-subquery/star). Fails on the not-yet-added SensitiveCheckResult.projections field, as intended. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_sensitive_projection_maskable.py | 235 ++++++++++++++++++++ 1 file changed, 235 insertions(+) create mode 100644 tests/test_sensitive_projection_maskable.py diff --git a/tests/test_sensitive_projection_maskable.py b/tests/test_sensitive_projection_maskable.py new file mode 100644 index 00000000..2465573d --- /dev/null +++ b/tests/test_sensitive_projection_maskable.py @@ -0,0 +1,235 @@ +"""ACE-041 slice 1: the sensitive-projection guard must, in addition to its existing +refuse decision, expose per-offending-projection *traceability* that a later masking slice +consumes — whether the projection is a 1:1 image of a single OUTPUT column (MASKABLE, with its +0-based output index) or has the sensitive value buried in an expression / function / scalar +subquery / star expansion (MUST-REFUSE). This slice adds NO runtime behaviour change: a sensitive +projection is still refused exactly as today; the maskable/index info is latent, exercised here. + +Fixture style mirrors tests/test_guard_context.py::_org (an inline Organization). Synthetic, +generic names only — `customers`, `orders`, `archived_customers`, `x` — no real data. +""" + +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() -> "m.Organization": + """One org with two sensitive columns on `customers` (`ssn`, `email`), a second sensitive-`ssn` + table (`archived_customers`) for set-operation arms, a non-sensitive `orders` for joins, and a + table `x` that carries a column literally named `ssn` which is NOT sensitive (the same-name + decoy).""" + customers = m.Table( + name="customers", schema="public", storage_connection="c", grain=["id"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="name", type="string"), + m.Column(name="dept", type="string"), + m.Column(name="ssn", type="string", sensitive=True), + 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="cust_id", type="integer"), + ], + ) + archived = m.Table( + name="archived_customers", schema="public", storage_connection="c", grain=["id"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="ssn", type="string", sensitive=True), + ], + ) + # `x.ssn` is a column literally named `ssn` on a DIFFERENT, non-sensitive table — the decoy the + # guard must NOT flag when it is the resolved source. + x = m.Table( + name="x", schema="public", storage_connection="c", grain=["ssn"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="ssn", type="string"), + ], + ) + sa = m.SubjectArea(name="area", description="d", tables_defined=[customers, orders, archived, x]) + return m.Organization(organization="AcmeCorp", version=1, subject_areas=[sa]) + + +def _classify(sql: str) -> tuple["rt.SensitiveCheckResult", list, list]: + """Run the guard and split its projections into (result, maskable[(index, column)], refuse[column]).""" + res = rt.check_sensitive_projection(sql, _org()) + maskable = sorted((p.output_index, p.column) for p in res.projections if p.maskable) + refuse = sorted(p.column for p in res.projections if not p.maskable) + return res, maskable, refuse + + +# --------------------------------------------------------------------------- +# Maskable: a bare col / t.col / either through a simple AS alias -> 1:1 output column. +# Each asserts the classification AND the exact 0-based output index. +# --------------------------------------------------------------------------- + +MASKABLE_CASES = [ + ("SELECT ssn FROM customers", [(0, "customers.ssn")]), + ("SELECT c.ssn FROM customers c", [(0, "customers.ssn")]), + ("SELECT c.ssn AS taxid FROM customers c", [(0, "customers.ssn")]), + ("SELECT id, ssn, name FROM customers", [(1, "customers.ssn")]), + ("SELECT ssn, email FROM customers", [(0, "customers.ssn"), (1, "customers.email")]), + ("SELECT ssn, ssn FROM customers", [(0, "customers.ssn"), (1, "customers.ssn")]), + ("SELECT c.ssn FROM customers c JOIN orders o ON o.cust_id = c.id", [(0, "customers.ssn")]), + ("SELECT ssn FROM customers ORDER BY ssn", [(0, "customers.ssn")]), + ("SELECT ssn FROM customers LIMIT 5", [(0, "customers.ssn")]), + # A derived table whose only real table in scope is `customers`, so the bare `ssn` resolves. + ("SELECT ssn FROM (SELECT ssn FROM customers) t", [(0, "customers.ssn")]), + # A CTE puts BOTH the cte name `t` and the physical `customers` into flat scope, so the bare + # column can't be pinned to one physical table -> the ref falls back to the bare name, but the + # PROJECTION is still a 1:1 image of output column 0, so it stays maskable at index 0. + ("WITH t AS (SELECT ssn FROM customers) SELECT ssn FROM t", [(0, "ssn")]), + # Every set-operation arm is walked; each arm's `ssn` is maskable at its own index 0. + ("SELECT ssn FROM customers UNION ALL SELECT ssn FROM archived_customers", + [(0, "archived_customers.ssn"), (0, "customers.ssn")]), + ("SELECT ssn FROM customers INTERSECT SELECT ssn FROM archived_customers", + [(0, "archived_customers.ssn"), (0, "customers.ssn")]), + ("SELECT ssn FROM customers EXCEPT SELECT ssn FROM archived_customers", + [(0, "archived_customers.ssn"), (0, "customers.ssn")]), +] + + +@pytest.mark.parametrize("sql, expected", MASKABLE_CASES) +def test_maskable_with_output_index(sql, expected): + res, maskable, refuse = _classify(sql) + # Runtime behaviour is unchanged: a sensitive projection is still refused. + assert res.action == "refuse", sql + # Nothing in these queries must-refuses — every offending projection is deterministically maskable. + assert refuse == [], sql + assert maskable == sorted(expected), sql + + +# --------------------------------------------------------------------------- +# Must-refuse: sensitive value buried in an expression / function / scalar subquery. +# These have NO deterministic 1:1 output column, so they are never maskable. +# --------------------------------------------------------------------------- + +MUST_REFUSE_CASES = [ + ("SELECT UPPER(ssn) FROM customers", ["customers.ssn"]), + ("SELECT ssn || '-x' FROM customers", ["customers.ssn"]), + ("SELECT SUBSTR(ssn, 1, 3) FROM customers", ["customers.ssn"]), + ("SELECT CASE WHEN id > 0 THEN ssn END FROM customers", ["customers.ssn"]), + ("SELECT COALESCE(ssn, '') FROM customers", ["customers.ssn"]), + ("SELECT (SELECT ssn FROM customers LIMIT 1) AS s", ["customers.ssn"]), + # A sensitive column wrapped in an expression in ONE arm of a set operation makes the WHOLE + # query must-refuse (no arm is maskable, so a later slice cannot mask its way to a safe answer). + ("SELECT UPPER(ssn) FROM customers UNION ALL SELECT id FROM customers", ["customers.ssn"]), +] + + +@pytest.mark.parametrize("sql, expected", MUST_REFUSE_CASES) +def test_must_refuse_never_maskable(sql, expected): + res, maskable, refuse = _classify(sql) + assert res.action == "refuse", sql + assert maskable == [], sql + assert refuse == sorted(expected), sql + + +def test_star_is_must_refuse_for_every_sensitive_column(): + """`SELECT *` expands a whole table; its output columns can't be pinned to a single + deterministic index, so each sensitive column is must-refuse (fail-closed), never maskable.""" + res, maskable, refuse = _classify("SELECT * FROM customers") + assert res.action == "refuse" + assert maskable == [] + assert refuse == ["customers.email", "customers.ssn"] + + +# --------------------------------------------------------------------------- +# No offending projection -> allow, and nothing to mask or refuse. +# --------------------------------------------------------------------------- + +ALLOW_CASES = [ + "SELECT COUNT(ssn) FROM customers", + "SELECT COUNT(DISTINCT ssn) FROM customers", + "SELECT dept, COUNT(*) FROM customers WHERE ssn = :x GROUP BY dept", + "SELECT COUNT(*) FROM customers c JOIN x ON c.ssn = x.ssn", + "SELECT name FROM customers ORDER BY ssn", +] + + +@pytest.mark.parametrize("sql", ALLOW_CASES) +def test_allowed_queries_report_nothing(sql): + res, maskable, refuse = _classify(sql) + assert res.action == "allow", sql + assert res.projections == [], sql + assert maskable == [] and refuse == [], sql + + +# --------------------------------------------------------------------------- +# Edge cases. +# --------------------------------------------------------------------------- + +def test_schema_qualified_projection_is_maskable(): + res, maskable, refuse = _classify("SELECT public.customers.ssn FROM public.customers") + assert res.action == "refuse" + assert refuse == [] + assert maskable == [(0, "customers.ssn")] + + +def test_same_named_nonsensitive_column_on_other_table_not_flagged(): + """`x.ssn` is literally named `ssn` but is NOT sensitive, so projecting it is allowed.""" + res, maskable, refuse = _classify("SELECT ssn FROM x") + assert res.action == "allow" + assert res.projections == [] + + +def test_quoted_case_variant_matches_current_behaviour(): + """DESIGN QUESTION (see report): sensitive-name matching is CASE-SENSITIVE today (unlike the + table/column-scope guards, which case-fold). A quoted/upper-case `"SSN"` therefore does NOT + match the lower-case declared `ssn`, so it is allowed. Hardening this to case-insensitive + matching would REFUSE more, i.e. change runtime behaviour, which is out of scope for this + slice — so this test pins the current behaviour rather than the (arguably) desired one.""" + res, maskable, refuse = _classify('SELECT "SSN" FROM customers') + assert res.action == "allow" + assert res.projections == [] + + +def test_unparseable_sql_reports_nothing_and_never_maskable(): + """An unparseable statement degrades to allow (the guard's documented posture) with an empty + projections list — so fail-closed's 'never maskable' holds (nothing is offered as maskable).""" + res, maskable, refuse = _classify("NOT SQL AT ALL ;;;") + assert res.action == "allow" + assert res.projections == [] + + +# --------------------------------------------------------------------------- +# Backwards-compat + ctx parity: the legacy fields are unchanged and the new +# maskable/index info is identical with or without a shared GuardContext. +# --------------------------------------------------------------------------- + +def test_legacy_fields_unchanged(): + res = rt.check_sensitive_projection("SELECT ssn FROM customers", _org()) + assert res.action == "refuse" + assert res.columns == ["customers.ssn"] + assert "customers.ssn" in res.reason + assert res.suggestion is not None + + +def test_maskable_index_is_ctx_invariant(): + org = _org() + sql = "SELECT id, ssn, name FROM customers" + ctx = rt.build_guard_context(sql, org) + no_ctx = rt.check_sensitive_projection(sql, org) + with_ctx = rt.check_sensitive_projection(sql, org, ctx=ctx) + assert no_ctx.as_dict() == with_ctx.as_dict() + assert [(p.output_index, p.column, p.maskable) for p in with_ctx.projections] == [ + (1, "customers.ssn", True) + ] From a9c8741d95578c88b8bbbd3bda58e4351d582716 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:05:06 -0700 Subject: [PATCH 02/11] ACE-041 slice1: check_sensitive_projection -> maskable + output index, tests green MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the sensitive-projection guard so each offending projection also carries the traceability a later masking slice needs, without changing runtime behaviour (the executor still refuses on ANY sensitive projection, exactly as today). - New frozen SensitiveProjection(column, maskable, output_index) and a SensitiveCheckResult.projections: list[SensitiveProjection] field (as_dict serializes it). Legacy action/columns/reason/suggestion are byte-identical. - Per projection (per set-operation arm), classify by SHAPE: (a) * / t.* -> must-refuse (no single deterministic output index), (b) col / t.col / AS x -> maskable at its 0-based projection index, (c) sensitive col buried in expr/func/scalar-subquery -> must-refuse. Fail-closed: only a bare/qualified/aliased single-column projection is maskable. - Reuses _output_selects (every arm), _resolve_col_table, _direct_from_tables, _count_protects, _sensitive_by_table, build_guard_context — single sqlglot parse, no second parser. Offending-set computation is preserved exactly; two small helpers (_sensitive_col_ref, _classify_projection) factor the shape logic. Caller execute_sql.py::_model_safety unchanged (reads action/reason/suggestion only). 33 new tests pass; full suite 1562 passed / 1 skipped, ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../agami-core/src/semantic_model/runtime.py | 137 +++++++++++++++--- 1 file changed, 117 insertions(+), 20 deletions(-) diff --git a/packages/agami-core/src/semantic_model/runtime.py b/packages/agami-core/src/semantic_model/runtime.py index 1553ba9c..e06b825f 100644 --- a/packages/agami-core/src/semantic_model/runtime.py +++ b/packages/agami-core/src/semantic_model/runtime.py @@ -424,12 +424,44 @@ def as_dict(self) -> dict[str, Any]: # --------------------------------------------------------------------------- +@dataclass(frozen=True) +class SensitiveProjection: + """One offending projection of a `sensitive` column, carrying the traceability a later masking + slice needs. `column` is the offending reference — "table.column" when the source table is + pinned, or a bare "column" when it cannot be (conservative fallback). `maskable` is True iff + this projection is a 1:1 image of a single OUTPUT column — a bare `col`, a table-qualified + `t.col`, or either through a simple `AS alias` — so a masker can deterministically replace that + one output column; False when the sensitive value is buried in an expression / function / + scalar subquery, or comes from a `*` expansion, where no single output column is a clean image + of the raw value (fail-closed → the query must be refused, never masked). `output_index` is the + 0-based position of this projection in its SELECT arm's projection list when `maskable`, else + None. Set-operation (UNION/INTERSECT/EXCEPT) arms are each walked, so one query can carry + several entries that share an `output_index` (one per arm).""" + + column: str + maskable: bool + output_index: Optional[int] = None + + def as_dict(self) -> dict[str, Any]: + return { + "column": self.column, + "maskable": self.maskable, + "output_index": self.output_index, + } + + @dataclass class SensitiveCheckResult: action: str # "allow" | "refuse" columns: list[str] = field(default_factory=list) # offending "table.column" / "column" reason: str = "" suggestion: Optional[str] = None + # Per-offending-projection traceability (ACE-041): the maskable-vs-must-refuse decision and, for + # maskable ones, the output-column index. Latent in THIS slice (the caller still refuses on any + # sensitive projection); a later masking slice reads it. Deterministic order: arm order, then + # projection index; duplicates and multi-arm collisions on `output_index` are preserved (unlike + # `columns`, which is the de-duplicated sorted set the refuse message is built from). + projections: list["SensitiveProjection"] = field(default_factory=list) def as_dict(self) -> dict[str, Any]: return { @@ -437,6 +469,7 @@ def as_dict(self) -> dict[str, Any]: "columns": self.columns, "reason": self.reason, "suggestion": self.suggestion, + "projections": [p.as_dict() for p in self.projections], } @@ -501,6 +534,81 @@ def _output_selects(node: "exp.Expression") -> list["exp.Select"]: return [] +def _sensitive_col_ref( + col: "exp.Column", + scope: dict[str, str], + by_table: dict[str, set[str]], + allnames: set[str], +) -> Optional[str]: + """The offending "table.column" (or bare "column") IF `col` projects a raw sensitive value, + else None. Byte-identical to the historical inline detection: a column whose name is sensitive + and is not COUNT-protected offends; it resolves to "table.column" when the table is pinned and + that table declares the column sensitive, to the bare name when the table is ambiguous + (conservative — sensitive somewhere in scope), and to None when it binds to a table that does + not declare it sensitive (a same-named column on a non-sensitive table).""" + if col.name not in allnames or _count_protects(col): + return None + tbl = _resolve_col_table(col, scope) + if tbl is None: + return col.name # ambiguous but sensitive somewhere in scope → conservative bare name + if tbl in by_table and col.name in by_table[tbl]: + return f"{tbl}.{col.name}" + return None # same-named column on a non-sensitive table → not offending + + +def _classify_projection( + proj: "exp.Expression", + index: int, + scope: dict[str, str], + direct: set[str], + by_table: dict[str, set[str]], + allnames: set[str], + offending: set[str], + projections: list["SensitiveProjection"], +) -> None: + """Classify ONE output projection at position `index` in its SELECT arm, appending any offending + sensitive column(s) to `offending` (the historical refuse set, computed byte-identically to + before) and one typed `SensitiveProjection` per offending column to `projections`. + + Three shapes, fail-closed by default (ACE-041): + (a) `*` / `t.*` — expand the starred table(s)' sensitive columns (exactly as the historical + star branch did) and mark each MUST-REFUSE: a star's output columns can't be pinned to a + single deterministic index. + (b) a bare `col`, a table-qualified `t.col`, or either through a simple `AS alias` — a 1:1 + image of one output column, so it is MASKABLE at `index`. + (c) anything else (function / arithmetic / CASE / COALESCE / scalar subquery / …) that still + contains a sensitive column — the raw value is buried, so MUST-REFUSE.""" + # (a) star projection — the historical `*` / `t.*` expansion, now fail-closed to must-refuse. + is_star = isinstance(proj, exp.Star) or ( + isinstance(proj, exp.Column) and isinstance(proj.this, exp.Star) + ) + if is_star: + qualifier = proj.table if isinstance(proj, exp.Column) else None + tables = {scope.get(qualifier, qualifier)} if qualifier else direct + for tbl in tables: + for c in sorted(by_table.get(tbl, set())): + ref = f"{tbl}.{c}" + offending.add(ref) + projections.append(SensitiveProjection(ref, maskable=False, output_index=None)) + return + + # (b) a single column projection (optionally behind a simple `AS` alias) → maskable at `index`. + inner = proj.this if isinstance(proj, exp.Alias) else proj + if isinstance(inner, exp.Column) and not isinstance(inner.this, exp.Star): + ref = _sensitive_col_ref(inner, scope, by_table, allnames) + if ref is not None: + offending.add(ref) + projections.append(SensitiveProjection(ref, maskable=True, output_index=index)) + return + + # (c) sensitive column buried in an expression / function / scalar subquery → must-refuse. + for col in inner.find_all(exp.Column): + ref = _sensitive_col_ref(col, scope, by_table, allnames) + if ref is not None: + offending.add(ref) + projections.append(SensitiveProjection(ref, maskable=False, output_index=None)) + + def check_sensitive_projection( sql: str, org: Organization, ctx: "GuardContext | None" = None ) -> SensitiveCheckResult: @@ -528,31 +636,19 @@ def check_sensitive_projection( if tree is None or tree.find(exp.Select) is None: return SensitiveCheckResult("allow") + # Walk every OUTPUT-bearing arm's projection list, classifying each projection into the shared + # `offending` set (the refuse decision, unchanged) and the per-projection `projections` list + # (the ACE-041 maskable/index traceability a later slice consumes). Index = position in THIS + # arm's projection list, which is the arm's output-column index. offending: set[str] = set() + projections: list[SensitiveProjection] = [] for sel in _output_selects(tree): scope = _tables_in_scope(sel) direct = _direct_from_tables(sel) - for proj in sel.expressions: - # (a) a raw projection of a sensitive column, not protected by COUNT - for col in proj.find_all(exp.Column): - if col.name not in allnames or _count_protects(col): - continue - tbl = _resolve_col_table(col, scope) - if tbl is None: - offending.add(col.name) # ambiguous + sensitive somewhere → conservative - elif tbl in by_table and col.name in by_table[tbl]: - offending.add(f"{tbl}.{col.name}") - # same-named column on a non-sensitive table → not offending - # (b) `*` / `t.*` that would expand a directly-FROM'd table holding sensitive cols - is_star = isinstance(proj, exp.Star) or ( - isinstance(proj, exp.Column) and isinstance(proj.this, exp.Star) + for index, proj in enumerate(sel.expressions): + _classify_projection( + proj, index, scope, direct, by_table, allnames, offending, projections ) - if is_star: - qualifier = proj.table if isinstance(proj, exp.Column) else None - tables = {scope.get(qualifier, qualifier)} if qualifier else direct - for tbl in tables: - for c in sorted(by_table.get(tbl, set())): - offending.add(f"{tbl}.{c}") if not offending: return SensitiveCheckResult("allow") @@ -565,6 +661,7 @@ def check_sensitive_projection( + " — sensitive columns may be counted or filtered, not output raw.", suggestion="Aggregate it (e.g. COUNT(DISTINCT )) for a count, or omit it and " "select the entity's non-sensitive key (e.g. id) instead.", + projections=projections, ) From 2564b9ad15cc3b61741fcc00c3467878eb95a746 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:13:33 -0700 Subject: [PATCH 03/11] ACE-041 slice2: failing policy tests Replace the data_protection stub regression pin with the real contract: a `provable` (deterministically maskable) verdict maps to `mask` in every tier; `heuristic`/`uncertain`/unexpected certainty fails closed to `reject` in every tier. Two maskable cases fail red against the current fail-closed stub, as intended (test-first). Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_guardrail.py | 55 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 5 deletions(-) diff --git a/tests/test_guardrail.py b/tests/test_guardrail.py index cda807c6..0d8fbe81 100644 --- a/tests/test_guardrail.py +++ b/tests/test_guardrail.py @@ -167,13 +167,58 @@ def test_policy_unknown_class_raises(): policy(bad) -# ── policy stubs for the data-protection / governance branches (regression pins) ─ +# ── policy: the data-protection branch masks when maskable, else fails closed ─ +# Data-protection ENFORCES in every tier (tier never downgrades it, unlike governance): the gate +# keys on `certainty` exactly the way safety does. `provable` means the gate traced every offending +# sensitive projection to a clean 1:1 image of a single output column, so a masker can +# deterministically replace it → `mask`. Anything less certain (`heuristic`/`uncertain`, or an +# unexpected value) means the raw value is buried / untraceable and cannot be safely masked, so it +# fails closed to `reject` — mirroring safety's "uncertainty ⇒ reject". -def test_policy_data_protection_stub_fails_closed(): - v = _safety(cls="data_protection", rule="sensitive_projection") - # Stub branch: fail-closed until the data-protection gates fill mask/row_filter selection. - assert policy(v) == "reject" +def _data_protection(**over) -> Verdict: + base = dict( + cls="data_protection", + rule="sensitive_projection", + severity="high", + certainty="provable", + detail="query projects raw values of sensitive column(s)", + remediation="aggregate, filter, or drop the sensitive column", + ) + base.update(over) + return Verdict(**base) + + +def test_policy_data_protection_maskable_masks_every_tier(): + # A `provable` data-protection verdict is deterministically maskable → `mask`, identically in + # every tier (the decision does not depend on tier — data-protection always enforces). + v = _data_protection(certainty="provable") + for tier in TIERS: + assert policy(v, tier) == "mask" + + +def test_policy_data_protection_unmaskable_rejects_every_tier(): + # `heuristic`/`uncertain` ⇒ the sensitive value cannot be traced to one output column, so it + # cannot be safely masked → fail closed to `reject`, identically in every tier. + for certainty in ("heuristic", "uncertain"): + v = _data_protection(certainty=certainty) + for tier in TIERS: + assert policy(v, tier) == "reject" + + +def test_policy_data_protection_fails_closed_on_edge_certainty(): + # Fail-closed default: ONLY `provable` masks; any other/unexpected certainty rejects. + v = _data_protection(certainty="bogus") + for tier in TIERS: + assert policy(v, tier) == "reject" + + +def test_policy_data_protection_default_tier_masks_when_maskable(): + # Default tier is oss; a maskable data-protection verdict still masks (tier is irrelevant here). + assert policy(_data_protection()) == "mask" + + +# ── policy stubs for the governance branch (regression pin) ────────────────── def test_policy_governance_stub_warns_by_default_enforces_at_enterprise(): From d103bec7e7cd5e48304a4dffaa28e12522bde1bc Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:16:19 -0700 Subject: [PATCH 04/11] ACE-041 slice2: data_protection policy branch (mask/reject), tests green MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement policy()'s data_protection branch. It keys on `certainty`, exactly like the safety branch: `provable` (the gate proved every offending sensitive projection is a deterministically maskable 1:1 image of one output column) -> `mask`; any weaker certainty (`heuristic`/`uncertain`/unexpected) -> `reject`, fail-closed, since the raw value cannot be safely masked. It enforces in every tier — tier is deliberately not consulted, unlike governance. Which output columns to mask travels separately on the guard result, not on the Verdict; policy() only decides mask-vs-reject. Re-vendored via `dev.py sync-lib`; plugins/agami/lib/guardrail.py is byte-identical to src (no unexpected drift). Guardrail tests pass against both src and the vendored lib; full suite green. Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/agami-core/src/guardrail.py | 19 ++++++++++++++----- plugins/agami/lib/guardrail.py | 19 ++++++++++++++----- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/packages/agami-core/src/guardrail.py b/packages/agami-core/src/guardrail.py index a0408862..4880d539 100644 --- a/packages/agami-core/src/guardrail.py +++ b/packages/agami-core/src/guardrail.py @@ -192,9 +192,15 @@ def policy(verdict: Verdict, tier: Tier = "oss") -> Action: - **safety** → always ``reject``, every tier; ``certainty == 'uncertain'`` also ⇒ ``reject`` (fail-closed on doubt). Fully implemented here. - - **data_protection** → ``mask`` / ``row_filter`` where the gate names one, else ``reject`` - (fail-closed). *Stub:* the data-protection gates fill this branch in later work; until then - it fails closed. + - **data_protection** → ``mask`` when the gate proves the sensitive projection is safely + maskable (``certainty == 'provable'``), else ``reject`` (fail-closed). It ENFORCES in every + tier — unlike governance, tier never downgrades a data-protection decision — and keys on + ``certainty`` exactly the way safety does: ``provable`` means the gate traced every offending + sensitive projection to a clean 1:1 image of a single output column, so a masker can + deterministically replace it; anything less certain (``heuristic`` / ``uncertain``, or an + unexpected value) means the raw value is buried / untraceable and cannot be safely masked, so + it fails closed. (Which output columns to mask travels separately on the guard result, not on + the ``Verdict`` — ``policy`` only decides mask-vs-reject.) - **governance** → graded by (severity, certainty, tier): ``reject``/``rewrite`` only for a ``provable`` verdict under an enforcing tier, else ``warn`` (OSS warns · SaaS recommends · Enterprise enforces). *Stub:* the governance gates fill this branch in later work. @@ -203,8 +209,11 @@ def policy(verdict: Verdict, tier: Tier = "oss") -> Action: # Safety is absolute: reject in every tier; uncertainty is already a reject (fail-closed). return "reject" if verdict.cls == "data_protection": - # Stub — fill mask/row_filter selection later. Fail closed until then (the safe default). - return "reject" + # Enforces in every tier (tier is deliberately not consulted here). A `provable` verdict is + # the gate's proof that every offending projection is a deterministically maskable 1:1 image + # of one output column → `mask`. Any weaker certainty means the sensitive value cannot be + # safely masked, so fail closed to `reject` — mirroring safety's "uncertainty ⇒ reject". + return "mask" if verdict.certainty == "provable" else "reject" if verdict.cls == "governance": # Stub — the graded (severity, certainty, tier) logic is filled in later. if tier == "enterprise" and verdict.certainty == "provable": diff --git a/plugins/agami/lib/guardrail.py b/plugins/agami/lib/guardrail.py index a0408862..4880d539 100644 --- a/plugins/agami/lib/guardrail.py +++ b/plugins/agami/lib/guardrail.py @@ -192,9 +192,15 @@ def policy(verdict: Verdict, tier: Tier = "oss") -> Action: - **safety** → always ``reject``, every tier; ``certainty == 'uncertain'`` also ⇒ ``reject`` (fail-closed on doubt). Fully implemented here. - - **data_protection** → ``mask`` / ``row_filter`` where the gate names one, else ``reject`` - (fail-closed). *Stub:* the data-protection gates fill this branch in later work; until then - it fails closed. + - **data_protection** → ``mask`` when the gate proves the sensitive projection is safely + maskable (``certainty == 'provable'``), else ``reject`` (fail-closed). It ENFORCES in every + tier — unlike governance, tier never downgrades a data-protection decision — and keys on + ``certainty`` exactly the way safety does: ``provable`` means the gate traced every offending + sensitive projection to a clean 1:1 image of a single output column, so a masker can + deterministically replace it; anything less certain (``heuristic`` / ``uncertain``, or an + unexpected value) means the raw value is buried / untraceable and cannot be safely masked, so + it fails closed. (Which output columns to mask travels separately on the guard result, not on + the ``Verdict`` — ``policy`` only decides mask-vs-reject.) - **governance** → graded by (severity, certainty, tier): ``reject``/``rewrite`` only for a ``provable`` verdict under an enforcing tier, else ``warn`` (OSS warns · SaaS recommends · Enterprise enforces). *Stub:* the governance gates fill this branch in later work. @@ -203,8 +209,11 @@ def policy(verdict: Verdict, tier: Tier = "oss") -> Action: # Safety is absolute: reject in every tier; uncertainty is already a reject (fail-closed). return "reject" if verdict.cls == "data_protection": - # Stub — fill mask/row_filter selection later. Fail closed until then (the safe default). - return "reject" + # Enforces in every tier (tier is deliberately not consulted here). A `provable` verdict is + # the gate's proof that every offending projection is a deterministically maskable 1:1 image + # of one output column → `mask`. Any weaker certainty means the sensitive value cannot be + # safely masked, so fail closed to `reject` — mirroring safety's "uncertainty ⇒ reject". + return "mask" if verdict.certainty == "provable" else "reject" if verdict.cls == "governance": # Stub — the graded (severity, certainty, tier) logic is filled in later. if tier == "enterprise" and verdict.certainty == "provable": From 1fdc7541365ace62f3bf63b4a7981e6611e7e967 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:31:41 -0700 Subject: [PATCH 05/11] ACE-041 slice3: failing mask e2e + executor + case-fold + unparseable tests Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_ace041_masking.py | 335 +++++++++++++++++++++++++++++++++++ 1 file changed, 335 insertions(+) create mode 100644 tests/test_ace041_masking.py diff --git a/tests/test_ace041_masking.py b/tests/test_ace041_masking.py new file mode 100644 index 00000000..9ac5d28b --- /dev/null +++ b/tests/test_ace041_masking.py @@ -0,0 +1,335 @@ +"""ACE-041 slice 3 — the behaviour flip: a *maskable* sensitive projection is no longer refused, +it is EXECUTED and the sensitive OUTPUT column is redacted to a token post-execution; an +*untraceable* sensitive projection still refuses (fail-closed). This pins the whole chain: + + guard classifies (runtime.check_sensitive_projection) -> _model_safety builds a data_protection + Verdict + runs guardrail.policy -> execute_guarded applies the mask plan to result.rows -> the + tool edge records `applied[{mask}]` on BOTH surfaces. + +Also covers the slice-3 case-fold fix (`SELECT SSN` / `SELECT "SSN"` now match a lower-case declared +`ssn`, consistent with the table/column-scope gates) and the unparseable fail-closed decision (this +branch has NO ACE-037 upstream unscopable gate, so the PII path itself refuses an unparseable +statement when the model declares PII). + +Synthetic, generic names only — `customers`, `orders`, `x`, `AcmeCorp`. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("pydantic") +pytest.importorskip("sqlglot") + +REPO_ROOT = Path(__file__).resolve().parent.parent +PKG_SRC = REPO_ROOT / "packages" / "agami-core" / "src" +if str(PKG_SRC) not in sys.path: + sys.path.insert(0, str(PKG_SRC)) +sys.path.insert(0, str(REPO_ROOT / "plugins" / "agami" / "scripts")) + +import execute_sql # noqa: E402 +from semantic_model import models as m # noqa: E402 +from semantic_model import runtime as rt # noqa: E402 + + +def _org() -> "m.Organization": + """`customers` has two sensitive columns (`ssn`, `email`); `orders` is non-sensitive for joins; + `x` carries a column literally named `ssn` that is NOT sensitive (the same-name decoy).""" + customers = m.Table( + name="customers", schema="public", storage_connection="c", grain=["id"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="name", type="string"), + m.Column(name="ssn", type="string", sensitive=True), + 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="cust_id", type="integer")], + ) + x = m.Table( + name="x", schema="public", storage_connection="c", grain=["ssn"], + columns=[m.Column(name="id", type="integer"), m.Column(name="ssn", type="string")], + ) + sa = m.SubjectArea(name="area", description="d", tables_defined=[customers, orders, x]) + return m.Organization(organization="AcmeCorp", version=1, subject_areas=[sa]) + + +class _SpyExecutor: + """Records calls and returns a canned ``ExecResult`` — so a test can assert the executor RAN and + what it returned was redacted downstream (not blocked).""" + + def __init__(self, result: "execute_sql.ExecResult"): + self.calls: list[tuple] = [] + self._result = result + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> "execute_sql.ExecResult": + self.calls.append((vetted_sql, creds, profile)) + return self._result + + +@pytest.fixture +def guarded(monkeypatch): + """Wire ``execute_guarded`` to run the REAL ``_model_safety`` over ``_org()`` with dummy creds, so + a test drives the whole guard->policy->redact chain with only the executor faked.""" + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + + def _run(sql: str, spy: "_SpyExecutor", **kw): + return execute_sql.execute_guarded(sql, "acme", None, executor=spy, **kw) + + return _run + + +def _R(columns, rows, truncated=False): + return execute_sql.ExecResult(columns=list(columns), rows=[tuple(r) for r in rows], truncated=truncated) + + +# --------------------------------------------------------------------------- +# execute_guarded: a maskable projection PROCEEDS and the value is redacted. +# --------------------------------------------------------------------------- + + +def test_maskable_projection_redacts_the_output_value_and_does_not_block(guarded): + spy = _SpyExecutor(_R(["ssn"], [("111-22-3333",), ("999-88-7777",)])) + result = guarded("SELECT ssn FROM customers", spy) + assert spy.calls, "the query must PROCEED to the executor, not be blocked" + assert result.rows == [(execute_sql.REDACTION_TOKEN,), (execute_sql.REDACTION_TOKEN,)] + assert result.columns == ["ssn"] + assert result.masked_columns == ("customers.ssn",) + + +def test_only_the_sensitive_output_column_is_redacted(guarded): + spy = _SpyExecutor(_R(["name", "ssn"], [("Alice", "111"), ("Bob", "222")])) + result = guarded("SELECT name, ssn FROM customers", spy) + # name (index 0, non-sensitive) is untouched; ssn (index 1) is redacted. + tok = execute_sql.REDACTION_TOKEN + assert result.rows == [("Alice", tok), ("Bob", tok)] + assert result.masked_columns == ("customers.ssn",) + + +def test_masking_holds_on_a_row_capped_result(guarded): + # The redaction runs on the ALREADY-capped rows and preserves the truncated flag (composes with + # the existing cap, does not fight it). + spy = _SpyExecutor(_R(["ssn"], [("a",), ("b",)], truncated=True)) + result = guarded("SELECT ssn FROM customers", spy) + tok = execute_sql.REDACTION_TOKEN + assert result.rows == [(tok,), (tok,)] + assert result.truncated is True + + +def test_two_sensitive_columns_both_redacted(guarded): + spy = _SpyExecutor(_R(["ssn", "email"], [("111", "a@x.io")])) + result = guarded("SELECT ssn, email FROM customers", spy) + tok = execute_sql.REDACTION_TOKEN + assert result.rows == [(tok, tok)] + assert result.masked_columns == ("customers.email", "customers.ssn") + + +def test_join_query_masks_the_sensitive_column_and_returns(guarded): + spy = _SpyExecutor(_R(["ssn"], [("111",)])) + result = guarded( + "SELECT c.ssn FROM customers c JOIN orders o ON o.cust_id = c.id", spy + ) + assert spy.calls + assert result.rows == [(execute_sql.REDACTION_TOKEN,)] + + +# --------------------------------------------------------------------------- +# Untraceable / star projections still REFUSE (fail-closed), executor never runs. +# --------------------------------------------------------------------------- + + +def test_untraceable_projection_still_refuses(guarded): + spy = _SpyExecutor(_R(["s"], [("x",)])) + with pytest.raises(execute_sql.GuardRefused) as ei: + guarded("SELECT UPPER(ssn) FROM customers", spy) + assert ei.value.refusal.kind == "sensitive_columns" + assert spy.calls == [] # buried value cannot be masked -> refused before the executor + + +def test_star_projection_still_refuses(guarded): + spy = _SpyExecutor(_R(["a"], [("x",)])) + with pytest.raises(execute_sql.GuardRefused) as ei: + guarded("SELECT * FROM customers", spy) + # SELECT * is caught by the star-ban safety gate FIRST (before the sensitive gate) — still refused. + assert spy.calls == [] + assert ei.value.refusal.kind in ("select_star", "sensitive_columns") + + +def test_partially_maskable_union_fails_closed(guarded): + # One arm buries ssn in UPPER(...) -> not every offending projection is maskable -> uncertain -> + # reject (fail closed). The executor never runs. + spy = _SpyExecutor(_R(["s"], [("x",)])) + with pytest.raises(execute_sql.GuardRefused) as ei: + guarded("SELECT UPPER(ssn) FROM customers UNION ALL SELECT id FROM customers", spy) + assert ei.value.refusal.kind == "sensitive_columns" + assert spy.calls == [] + + +# --------------------------------------------------------------------------- +# Allowed queries are not masked at all (COUNT / WHERE / non-sensitive). +# --------------------------------------------------------------------------- + + +def test_count_of_sensitive_is_not_masked(guarded): + spy = _SpyExecutor(_R(["n"], [(5,)])) + result = guarded("SELECT COUNT(ssn) FROM customers", spy) + assert result.rows == [(5,)] # untouched + assert result.masked_columns == () + + +def test_nonsensitive_projection_is_not_masked(guarded): + spy = _SpyExecutor(_R(["name"], [("Alice",), ("Bob",)])) + result = guarded("SELECT name FROM customers WHERE ssn = :x", spy) + assert result.rows == [("Alice",), ("Bob",)] + assert result.masked_columns == () + + +# --------------------------------------------------------------------------- +# Case-fold fix (ACE-041 slice 3): the sensitive-name match is now case-insensitive. +# --------------------------------------------------------------------------- + + +def test_uppercase_unquoted_sensitive_column_is_now_masked(guarded): + spy = _SpyExecutor(_R(["SSN"], [("111",)])) + result = guarded("SELECT SSN FROM customers", spy) + assert spy.calls # proceeds + assert result.rows == [(execute_sql.REDACTION_TOKEN,)] + assert result.masked_columns == ("customers.ssn",) + + +def test_quoted_uppercase_sensitive_column_is_now_masked(guarded): + spy = _SpyExecutor(_R(["SSN"], [("111",)])) + result = guarded('SELECT "SSN" FROM customers', spy) + assert result.rows == [(execute_sql.REDACTION_TOKEN,)] + assert result.masked_columns == ("customers.ssn",) + + +def test_mixed_case_sensitive_column_is_now_masked(guarded): + spy = _SpyExecutor(_R(["Ssn"], [("111",)])) + result = guarded("SELECT Ssn FROM customers", spy) + assert result.rows == [(execute_sql.REDACTION_TOKEN,)] + assert result.masked_columns == ("customers.ssn",) + + +def test_same_named_nonsensitive_column_on_other_table_still_not_flagged(guarded): + # `x.ssn` is literally named `ssn` but NOT sensitive on `x` — projecting it (any case) is allowed. + spy = _SpyExecutor(_R(["ssn"], [("v",)])) + result = guarded("SELECT SSN FROM x", spy) + assert result.rows == [("v",)] # untouched + assert result.masked_columns == () + + +# --------------------------------------------------------------------------- +# Unparseable fail-closed (this branch has no upstream ACE-037 unscopable gate): +# an unparseable statement + a model that declares PII is REFUSED, never run. +# --------------------------------------------------------------------------- + + +def test_unparseable_with_declared_pii_fails_closed(guarded, monkeypatch): + # Force sqlglot to fail to parse (tree=None) while it is otherwise available, so we exercise the + # PARSE-FAILURE fail-closed branch (distinct from sqlglot-unavailable, which stays a fail-open). + monkeypatch.setattr(rt, "_parse_sql", lambda sql: None) + spy = _SpyExecutor(_R(["ssn"], [("111",)])) + with pytest.raises(execute_sql.GuardRefused) as ei: + guarded("SELECT ssn FROM customers", spy) + assert ei.value.refusal.kind == "sensitive_columns" + assert spy.calls == [] + + +# --------------------------------------------------------------------------- +# _model_safety returns a typed (sql, Refusal|None, MaskPlan|None). +# --------------------------------------------------------------------------- + + +def test_model_safety_returns_a_mask_plan_for_a_maskable_projection(monkeypatch): + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + sql, refusal, plan = execute_sql._model_safety("SELECT id, ssn FROM customers", "acme", None) + assert refusal is None + assert plan is not None + assert plan.indices == (1,) + assert plan.columns == ("customers.ssn",) + assert plan.token == execute_sql.REDACTION_TOKEN + + +def test_model_safety_rejects_an_untraceable_projection(monkeypatch): + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + sql, refusal, plan = execute_sql._model_safety("SELECT UPPER(ssn) FROM customers", "acme", None) + assert plan is None + assert refusal is not None and refusal.kind == "sensitive_columns" + + +def test_model_safety_allows_a_clean_query_with_no_plan(monkeypatch): + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + sql, refusal, plan = execute_sql._model_safety("SELECT name FROM customers", "acme", None) + assert refusal is None and plan is None + + +# --------------------------------------------------------------------------- +# The subprocess metadata channel: _emit_result_csv flags masked columns on stderr, and +# tools._executor_masked parses that line. +# --------------------------------------------------------------------------- + + +def test_emit_result_csv_flags_masked_columns_on_stderr(capsys): + result = execute_sql.ExecResult( + columns=["ssn"], rows=[("***",)], truncated=False, masked_columns=("customers.ssn",) + ) + execute_sql._emit_result_csv(result) + err = capsys.readouterr().err + lines = [json.loads(x) for x in err.splitlines() if x.strip().startswith("{")] + masked = [d["masked"] for d in lines if "masked" in d] + assert masked == [["customers.ssn"]] + + +def test_tools_executor_masked_parses_the_stderr_marker(): + import tools + + stderr = 'some notice\n{"masked": ["customers.ssn", "customers.email"]}\n' + assert tools._executor_masked(stderr) == ["customers.ssn", "customers.email"] + assert tools._executor_masked("nothing here") == [] + + +# --------------------------------------------------------------------------- +# tool_execute_sql (in-process, injected executor): masked VALUES in `data`, `applied[{mask}]` note. +# --------------------------------------------------------------------------- + + +def test_tool_execute_sql_masks_value_and_records_applied(monkeypatch): + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + + spy = _SpyExecutor(_R(["name", "ssn"], [("Alice", "111-22-3333")])) + tools.set_injected_executor(spy) + out = json.loads(tools.tool_execute_sql({"sql": "SELECT name, ssn FROM customers", "datasource": "acme"})) + + assert out["status"] == "ok" # NOT refused — the query ran + assert spy.calls # the injected executor ran behind the guard + assert out["data"]["rows"] == [["Alice", execute_sql.REDACTION_TOKEN]] # ssn redacted, name intact + assert {"mask": "customers.ssn"} in out["applied"] + + +def test_tool_execute_sql_untraceable_pii_is_refused_no_data(monkeypatch): + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + + spy = _SpyExecutor(_R(["s"], [("x",)])) + tools.set_injected_executor(spy) + out = json.loads(tools.tool_execute_sql({"sql": "SELECT UPPER(ssn) FROM customers", "datasource": "acme"})) + + assert out["status"] == "refused" + assert out["refusal"]["kind"] == "sensitive_columns" + assert "data" not in out + assert spy.calls == [] From 4bd39bb620f311e1271fdcdcdc83c782a04c84e5 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:33:19 -0700 Subject: [PATCH 06/11] ACE-041 slice3: case-insensitive sensitive-name match (close SELECT "SSN" bypass) _sensitive_by_table now folds table+column keys and preserves the declared spelling in the ref value; _sensitive_col_ref and the star branch fold-match. Flips the (now spec-obsolete) case-sensitive pin to assert mask, adds SSN/ mixed-case/quoted matrix rows + the non-sensitive same-name decoy. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../agami-core/src/semantic_model/runtime.py | 49 ++++++++++++------- tests/test_sensitive_projection_maskable.py | 30 +++++++++--- 2 files changed, 53 insertions(+), 26 deletions(-) diff --git a/packages/agami-core/src/semantic_model/runtime.py b/packages/agami-core/src/semantic_model/runtime.py index e06b825f..f80fe406 100644 --- a/packages/agami-core/src/semantic_model/runtime.py +++ b/packages/agami-core/src/semantic_model/runtime.py @@ -80,7 +80,7 @@ class GuardContext: 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]]" + sensitive_by_table: "tuple[dict[str, dict[str, str]], set[str]]" model_table_index: "dict[str, tuple]" @@ -473,16 +473,23 @@ def as_dict(self) -> dict[str, Any]: } -def _sensitive_by_table(org: Organization) -> tuple[dict[str, set[str]], set[str]]: - """(table name -> {sensitive column names}, union of all sensitive column names).""" - by_table: dict[str, set[str]] = {} +def _sensitive_by_table(org: Organization) -> tuple[dict[str, dict[str, str]], set[str]]: + """(lower(table) -> {lower(col): "DeclaredTable.DeclaredCol"}, {lower(col) for every sensitive col}). + + Keys are case-FOLDED so the sensitive-name match is case-insensitive — consistent with + `check_table_scope` / `check_column_scope`, which already fold (unquoted identifiers fold case in + Postgres and friends). Before ACE-041 slice 3 this index was case-sensitive, so `SELECT "SSN"` / + `SELECT SSN` slipped past a lower-case declared `ssn`; folding closes that PII bypass (it + masks/refuses MORE). The value preserves the DECLARED spelling so an offending ref names the + model's column regardless of how the query happened to spell it.""" + by_table: dict[str, dict[str, str]] = {} allnames: set[str] = set() for sa in org.subject_areas: for t in sa.tables_defined: for c in t.columns: if getattr(c, "sensitive", False): - by_table.setdefault(t.name, set()).add(c.name) - allnames.add(c.name) + by_table.setdefault(t.name.lower(), {})[c.name.lower()] = f"{t.name}.{c.name}" + allnames.add(c.name.lower()) return by_table, allnames @@ -537,22 +544,25 @@ def _output_selects(node: "exp.Expression") -> list["exp.Select"]: def _sensitive_col_ref( col: "exp.Column", scope: dict[str, str], - by_table: dict[str, set[str]], + by_table: dict[str, dict[str, str]], allnames: set[str], ) -> Optional[str]: """The offending "table.column" (or bare "column") IF `col` projects a raw sensitive value, - else None. Byte-identical to the historical inline detection: a column whose name is sensitive - and is not COUNT-protected offends; it resolves to "table.column" when the table is pinned and - that table declares the column sensitive, to the bare name when the table is ambiguous - (conservative — sensitive somewhere in scope), and to None when it binds to a table that does - not declare it sensitive (a same-named column on a non-sensitive table).""" - if col.name not in allnames or _count_protects(col): + else None. Matching is CASE-INSENSITIVE (ACE-041 slice 3), consistent with the table/column-scope + gates: a column whose folded name is sensitive and is not COUNT-protected offends; it resolves to + the DECLARED "table.column" when the table is pinned and that table declares the column sensitive, + to the folded bare name when the table is ambiguous (conservative — sensitive somewhere in scope), + and to None when it binds to a table that does not declare it sensitive (a same-named column on a + non-sensitive table).""" + cname = col.name.lower() + if cname not in allnames or _count_protects(col): return None tbl = _resolve_col_table(col, scope) if tbl is None: - return col.name # ambiguous but sensitive somewhere in scope → conservative bare name - if tbl in by_table and col.name in by_table[tbl]: - return f"{tbl}.{col.name}" + return cname # ambiguous but sensitive somewhere in scope → conservative folded bare name + tcols = by_table.get(tbl.lower()) + if tcols is not None and cname in tcols: + return tcols[cname] # the DECLARED "table.column" ref (query casing does not leak into it) return None # same-named column on a non-sensitive table → not offending @@ -561,7 +571,7 @@ def _classify_projection( index: int, scope: dict[str, str], direct: set[str], - by_table: dict[str, set[str]], + by_table: dict[str, dict[str, str]], allnames: set[str], offending: set[str], projections: list["SensitiveProjection"], @@ -586,8 +596,9 @@ def _classify_projection( qualifier = proj.table if isinstance(proj, exp.Column) else None tables = {scope.get(qualifier, qualifier)} if qualifier else direct for tbl in tables: - for c in sorted(by_table.get(tbl, set())): - ref = f"{tbl}.{c}" + tcols = by_table.get((tbl or "").lower(), {}) + for ckey in sorted(tcols): # folded key order → the historical sorted-by-name order + ref = tcols[ckey] # DECLARED "table.column" offending.add(ref) projections.append(SensitiveProjection(ref, maskable=False, output_index=None)) return diff --git a/tests/test_sensitive_projection_maskable.py b/tests/test_sensitive_projection_maskable.py index 2465573d..a8bfa878 100644 --- a/tests/test_sensitive_projection_maskable.py +++ b/tests/test_sensitive_projection_maskable.py @@ -191,13 +191,29 @@ def test_same_named_nonsensitive_column_on_other_table_not_flagged(): assert res.projections == [] -def test_quoted_case_variant_matches_current_behaviour(): - """DESIGN QUESTION (see report): sensitive-name matching is CASE-SENSITIVE today (unlike the - table/column-scope guards, which case-fold). A quoted/upper-case `"SSN"` therefore does NOT - match the lower-case declared `ssn`, so it is allowed. Hardening this to case-insensitive - matching would REFUSE more, i.e. change runtime behaviour, which is out of scope for this - slice — so this test pins the current behaviour rather than the (arguably) desired one.""" - res, maskable, refuse = _classify('SELECT "SSN" FROM customers') +@pytest.mark.parametrize("sql", [ + 'SELECT "SSN" FROM customers', # quoted upper-case + "SELECT SSN FROM customers", # unquoted upper-case + "SELECT Ssn FROM customers", # mixed case + "SELECT c.SSN FROM customers c", # qualified upper-case +]) +def test_case_insensitive_sensitive_name_is_now_masked(sql): + """ACE-041 slice 3 (behaviour CHANGE — was `allow` before this slice): the sensitive-name match + is now CASE-INSENSITIVE, consistent with the table/column-scope gates (which fold). An upper / + mixed / quoted spelling of a lower-case declared `ssn` therefore matches and is a maskable + projection at output index 0 — this masks/refuses MORE, closing the `SELECT "SSN"` PII bypass. + The declared ref (`customers.ssn`) names the model column regardless of the query's casing.""" + res, maskable, refuse = _classify(sql) + assert res.action == "refuse", sql + assert refuse == [], sql + assert maskable == [(0, "customers.ssn")], sql + + +def test_case_insensitive_decoy_on_nonsensitive_table_still_not_flagged(): + """The case-fold must not over-reach: `x.ssn` is a same-named column on a NON-sensitive table, so + an upper-case `SELECT SSN FROM x` is still allowed (it binds to `x`, which does not declare it + sensitive).""" + res, maskable, refuse = _classify("SELECT SSN FROM x") assert res.action == "allow" assert res.projections == [] From 5a7f74a249e01c169ceda7e4ec6418031622a634 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:37:11 -0700 Subject: [PATCH 07/11] ACE-041 slice3: redact at execute_guarded + mask plan wiring + applied[{mask}] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _model_safety now returns (sql, Refusal|None, MaskPlan|None): a maskable sensitive projection routes through a data_protection Verdict + policy -> MASK plan (query proceeds), an untraceable one fails closed. execute_guarded redacts result.rows at the single shared post-execute point (covers _run_ and any injected ports.Executor). ExecResult.masked_columns + a {"masked": …} stderr marker carry the applied[{mask}] note on both surfaces. Unparseable SQL with declared PII fails closed (no ACE-037 upstream gate on this branch). Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/agami-core/src/execute_sql.py | 167 +++++++++++++++++++++---- packages/agami-core/src/tools.py | 52 ++++++-- 2 files changed, 183 insertions(+), 36 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 341e2ca1..f48fbb2d 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -53,7 +53,7 @@ from typing import TYPE_CHECKING, Any import agami_paths -from guardrail import Refusal, Verdict +from guardrail import Refusal, Verdict, policy if TYPE_CHECKING: # ``Executor`` is the 5th port; imported only for type-checkers. At runtime ``execute_sql`` never @@ -113,6 +113,30 @@ class ExecResult: columns: list[str] rows: list[tuple] truncated: bool = False + # ACE-041: the DECLARED "table.column" refs of any OUTPUT columns redacted to REDACTION_TOKEN by + # the data-protection mask plan (empty for an unmasked result). The redacted VALUES already live + # in `rows`; this carries the `applied[{mask}]` note both surfaces record. Immutable default `()`. + masked_columns: tuple[str, ...] = () + + +# ACE-041: the string a masked sensitive value is replaced with in every row of a masked output +# column. A single named constant so the token is defined once for the redactor and the tests. +# ACE-041: token pending author confirm +REDACTION_TOKEN = "***" + + +@dataclass(frozen=True) +class MaskPlan: + """The post-execution redaction the data-protection gate proved safe: which OUTPUT column + positions to redact (`indices`, 0-based, aligned with the result's columns and every set-op arm) + and the DECLARED "table.column" refs those positions carry (`columns`, for the `applied[{mask}]` + note). Built by ``_model_safety`` only when ``guardrail.policy`` returns ``mask`` (every offending + projection is a deterministically-maskable 1:1 image of one output column), carried to + ``execute_guarded`` which applies it to ``result.rows``. Frozen — one plan per guarded call.""" + + indices: tuple[int, ...] + columns: tuple[str, ...] + token: str = REDACTION_TOKEN class ExecutorError(Exception): @@ -837,11 +861,21 @@ def _collect_cursor(cur: Any) -> ExecResult: return ExecResult(columns=columns, rows=[tuple(r) for r in fetched[:cap]], truncated=truncated) +def _flag_masked(columns: tuple[str, ...]) -> None: + """Signal which OUTPUT columns were PII-redacted to the subprocess caller — a non-error + ``{"masked": [...]}`` marker on stderr, the SAME metadata channel as ``{"truncated": …}`` (not a + new protocol). The redacted VALUES already ride the CSV on stdout; this line carries only the + ``applied[{mask}]`` note the parent records on the Envelope (ACE-041 slice 3).""" + sys.stderr.write(json.dumps({"masked": list(columns)}) + "\n") + + def _emit_result_csv(result: ExecResult) -> None: """Serialize an ``ExecResult`` to stdout as CSV — the subprocess/CLI wire. Byte-for-byte what the old inline cursor→CSV writer produced: header row then data rows, and a truncation marker on stderr when capped. This is the *single, final* text serialization for the fork path; the - in-process path skips it and returns the native rows straight to the tool edge.""" + in-process path skips it and returns the native rows straight to the tool edge. When the rows were + PII-masked, the redacted VALUES are already in ``result.rows`` (masked in ``execute_guarded`` + before this call), so the CSV carries them; a ``{"masked": …}`` stderr marker names the columns.""" if not result.columns: # cursor had no description → wrote nothing (e.g. a non-row statement) return writer = csv.writer(sys.stdout) @@ -850,6 +884,8 @@ def _emit_result_csv(result: ExecResult) -> None: writer.writerow(row) if result.truncated: _flag_truncated(_resolve_row_cap()) + if result.masked_columns: + _flag_masked(result.masked_columns) def _write_cursor_csv(cur: Any) -> None: @@ -920,15 +956,23 @@ def _refusal_from_verdict(kind: str, verdict: Verdict) -> Refusal: return Refusal(kind=kind, reason=verdict.detail, remediation=verdict.remediation) -def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusal | None]: - """Semantic-model safety pass before execution: fan-trap / chasm-trap pre-flight - + default_filters auto-application, over a model resolved from the DB (hosted) or disk (local). +def _model_safety( + sql: str, profile: str, area: str | None +) -> tuple[str, Refusal | None, MaskPlan | None]: + """Semantic-model safety pass before execution: fan-trap / chasm-trap pre-flight + scope + PII + + default_filters auto-application, over a model resolved from the DB (hosted) or disk (local). - Returns (sql_to_run, refusal). ``refusal`` is None to continue, or a typed ``Refusal`` the caller - (``execute_guarded``) raises as a ``GuardRefused`` — so both the subprocess and in-process paths - build the same envelope from it. Inert (returns the SQL unchanged) when the model package isn't - importable, or — on the LOCAL path only — when there is no model yet. On the HOSTED path a model - that can't be resolved fails closed (refuses), never runs unguarded (ACE-051).""" + Returns ``(sql_to_run, refusal, mask_plan)``: + - ``refusal`` is None to continue, or a typed ``Refusal`` the caller (``execute_guarded``) + raises as a ``GuardRefused`` — so both the subprocess and in-process paths build the same + envelope from it. + - ``mask_plan`` (ACE-041 slice 3) is None unless the PII gate proved a *maskable* sensitive + projection: the query then PROCEEDS and ``execute_guarded`` redacts the named output columns + post-execution. An *untraceable* sensitive projection sets ``refusal`` instead (fail-closed). + + Inert (SQL unchanged, no refusal, no plan) when the model package isn't importable, or — on the + LOCAL path only — when there is no model yet. On the HOSTED path a model that can't be resolved + fails closed (refuses), never runs unguarded (ACE-051).""" try: from semantic_model import runtime as RT except Exception: @@ -942,8 +986,8 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa "model_unavailable", "semantic-model package not importable; refusing to run unguarded on the " "hosted server", - ) - return sql, None # local: model package not available -> no-op + ), None + return sql, None, None # local: model package not available -> no-op org = _resolve_guard_model(profile) if org is None: @@ -954,8 +998,8 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa "model_unavailable", "no semantic model could be resolved (checked DB and disk); refusing to run " "unguarded on the hosted server", - ) - return sql, None # local: no model yet -> no-op (unchanged) + ), None + return sql, None, None # local: no model yet -> no-op (unchanged) # 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 @@ -968,41 +1012,79 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa # so the fan/chasm and sensitive checks below only evaluate in-scope tables. ts = RT.check_table_scope(sql, org, ctx=ctx) if ts is not None: - return sql, _refusal_from_verdict("table_out_of_scope", ts) + return sql, _refusal_from_verdict("table_out_of_scope", ts), 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, ctx=ctx) if star is not None: - return sql, _refusal_from_verdict("select_star", star) + return sql, _refusal_from_verdict("select_star", star), 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, ctx=ctx) if cs is not None: - return sql, _refusal_from_verdict("column_out_of_scope", cs) + return sql, _refusal_from_verdict("column_out_of_scope", cs), None pf = RT.pre_flight_check(sql, org, ctx=ctx) if pf.risk and pf.action == "refuse": - return sql, Refusal("preflight_refused", pf.reason, pf.suggestion or "") + return sql, Refusal("preflight_refused", pf.reason, pf.suggestion or ""), 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. + # Sensitive-column (PII) fail-closed on UNPARSEABLE SQL when the model declares PII (ACE-041 + # slice 3): a statement sqlglot cannot parse cannot be proven free of a raw sensitive projection. + # Unlike the other model guards (which degrade to allow — that fail-open is ACE-037's to close; + # ACE-037 is NOT on this branch), the PII path refuses a parse FAILURE outright rather than run a + # possibly-PII-projecting query. Only a parse failure with sqlglot PRESENT (ctx is not None and + # ctx.tree is None) is closed here; sqlglot-unavailable (ctx is None) stays the ACE-037-owned + # fail-open, and a model with no sensitive columns has nothing to protect. + if ctx is not None and ctx.tree is None and ctx.sensitive_by_table[1]: + return sql, Refusal( + "sensitive_columns", + "the query could not be parsed to verify it does not expose sensitive columns, so it " + "was refused.", + "Simplify the SQL so it parses, then retry.", + ), None + + # Sensitive-column (PII) guard — the deterministic chokepoint that protects PII identically for + # the agami-query skill, the local MCP server, and cron. A sensitive column may still be counted, + # filtered, grouped, or joined; only PROJECTING its raw value is offending. When it is, route the + # gate's verdict through the shared data-protection policy: a *maskable* projection (every + # offending projection is a clean 1:1 image of one output column) becomes a MASK plan and the + # query PROCEEDS with those columns redacted post-execution; an *untraceable* projection fails + # closed to a refusal (mirroring safety's "uncertainty ⇒ reject"). + mask_plan: MaskPlan | None = None sens = RT.check_sensitive_projection(sql, org, ctx=ctx) if sens.action == "refuse": - return sql, Refusal("sensitive_columns", sens.reason, sens.suggestion or "") + all_maskable = bool(sens.projections) and all( + p.maskable and p.output_index is not None for p in sens.projections + ) + verdict = Verdict( + cls="data_protection", + rule="sensitive_projection", + severity="high", + certainty="provable" if all_maskable else "uncertain", + detail=sens.reason, + remediation=sens.suggestion or "", + ) + # data_protection enforces in every tier (tier never downgrades it), so policy's default + # tier is correct here: provable → "mask", anything weaker → "reject" (fail closed). + if policy(verdict) == "mask": + indices = tuple(sorted({p.output_index for p in sens.projections})) + mask_plan = MaskPlan(indices=indices, columns=tuple(sens.columns)) + else: + return sql, Refusal("sensitive_columns", sens.reason, sens.suggestion or ""), None new_sql, applied = RT.apply_default_filters(sql, org, area=area, ctx=ctx) if applied: + # default_filters only add WHERE conditions, never change the projection list, so the mask + # plan's output indices stay valid against the rewritten SQL's result. sys.stderr.write(f"[agami] applied default_filters: {applied}\n") sql = new_sql - return sql, None + return sql, None, mask_plan # --------------------------------------------------------------------------- @@ -1143,18 +1225,49 @@ def execute_guarded( SQL both guards have passed. Raises ``GuardRefused`` carrying the typed ``Refusal`` on a refusal, and ``ExecutorError`` on a connect/run failure — so the subprocess ``main`` and the in-process MCP handler apply the same guard and build the same envelope. The row cap rides the request-scoped - ``_max_rows_override`` ContextVar the caller sets.""" + ``_max_rows_override`` ContextVar the caller sets. + + ACE-041: when the model-safety pass returns a ``MaskPlan`` (a *maskable* sensitive projection), + the query runs and the named output columns are redacted here, at the single shared post-execute + point — so masking covers every ``_run_`` AND any injected ``ports.Executor``, and composes + with the already-applied row cap (it redacts the capped rows in place).""" import sql_guard verdict = sql_guard.check_read_only(sql) if verdict is not None: raise GuardRefused(_refusal_from_verdict("permission", verdict), code=1) + mask_plan: MaskPlan | None = None if not no_safety: - sql, refusal = _model_safety(sql, profile, area) + sql, refusal, mask_plan = _model_safety(sql, profile, area) if refusal is not None: raise GuardRefused(refusal, code=1) creds = _load_credentials(profile) - return executor.execute(sql, creds, profile=profile) + result = executor.execute(sql, creds, profile=profile) + if mask_plan is not None: + result = _apply_mask_plan(result, mask_plan) + return result + + +def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: + """Redact the sensitive OUTPUT columns of an executor result post-execution — the SINGLE shared + masking point for EVERY executor (built-in ``_run_`` and an injected ``ports.Executor``). + Only the values at the plan's output indices become the token; every other value is untouched, + and the redacted rows plus the masked column refs are returned in a NEW ``ExecResult`` (it is + frozen). Runs on the already-row-capped rows, so it composes with the fetch bound rather than + fighting it. An index outside a row's width simply doesn't match (fail-safe — the value isn't + there to leak); by construction an output index equals its result-column position, so this holds.""" + if not plan.indices: + return result + mask_set = set(plan.indices) + redacted_rows = [ + tuple(plan.token if i in mask_set else v for i, v in enumerate(row)) for row in result.rows + ] + return ExecResult( + columns=result.columns, + rows=redacted_rows, + truncated=result.truncated, + masked_columns=plan.columns, + ) def main() -> int: diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index f1e28b71..b1232bde 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -880,6 +880,23 @@ def _executor_truncated(stderr: str | None) -> bool: return False +def _executor_masked(stderr: str | None) -> list[str]: + """The DECLARED "table.column" refs of the columns execute_sql PII-redacted (ACE-041). The + executor emits a non-error ``{"masked": [...]}`` line on stderr — the same metadata channel as + ``{"truncated": …}`` — AFTER redacting the VALUES on stdout, so the subprocess surface records the + same ``applied[{mask}]`` note the in-process surface does. Returns [] when nothing was masked.""" + for line in (stderr or "").splitlines(): + line = line.strip() + if line.startswith("{") and '"masked"' in line: + try: + masked = json.loads(line).get("masked") + except ValueError: + continue + if isinstance(masked, list): + return [str(c) for c in masked] + return [] + + # The composition-root executor (AH-012). ``None`` (the default) means "fork the execute_sql # subprocess" — the byte-identical local/single-user path. A consumer injects a ``ports.Executor`` # via ``create_app(adapters=…)`` to run execution IN-PROCESS behind the same guard (no fork, native @@ -907,11 +924,16 @@ def set_injected_executor(executor: Any | None) -> None: def _finalize_execution( columns: list, data_rows: list, truncated: bool, *, profile: str, sql: str, execution_ms: int, args: dict[str, Any], audit_id: str, + masked_columns: list[str] | None = None, ) -> str: """Shape a successful result (units + exact-render markdown + trust receipt) into the ok ``Envelope``, record it through the guardrail-audit chokepoint, and return the Envelope JSON. Shared by BOTH execution paths — the subprocess fork and the in-process executor — so a - successful query returns the identical envelope whichever ran it.""" + successful query returns the identical envelope whichever ran it. + + ``masked_columns`` (ACE-041) are the DECLARED "table.column" refs the data-protection gate + redacted in ``data_rows``; each is recorded as an ``applied[{mask}]`` note so the trust surface + reports what was masked (the redacted VALUES already ride ``data_rows``).""" # Deterministic, exact rendering — so the numbers a user verifies don't depend on # how the host LLM chooses to format them. `markdown` is the table to display # verbatim; `rows` stays raw (exact CSV values) for charting / programmatic use. @@ -938,8 +960,10 @@ def _finalize_execution( # (offer to approve/correct via the save_correction tool). "receipt": _resolve_receipt(profile, sql), } - # A bounded (transfer-capped) result is a valid answer that is flagged, not a refusal. + # A bounded (transfer-capped) result is a valid answer that is flagged, not a refusal. PII + # redaction is likewise a valid, flagged answer: one `{mask: "table.col"}` per redacted column. applied = [{"row_cap": len(data_rows)}] if truncated else [] + applied += [{"mask": col} for col in (masked_columns or [])] # Log the execution through the single chokepoint: the DB sink when AGAMI_DB_URL is set (one # query_executions row), else the local jsonl the skills use. Best-effort either way. @@ -963,12 +987,16 @@ def _finalize_execution( def _run_in_process( sql: str, profile: str, area: str | None, max_rows: int | None, executor: Any -) -> tuple[list, list, bool] | Refusal: +) -> tuple[list, list, bool, list] | Refusal: """Run through the in-process executor behind the shared guarded envelope (no subprocess, no CSV - round-trip). Returns ``(columns, data_rows, truncated)`` on success, or the typed ``Refusal`` on a - guard refusal / execution failure — the SAME refusal the subprocess path produces (``GuardRefused`` - now carries the typed ``Refusal``, incl. the model-safety detail), so the two paths build an - identical refused Envelope. + round-trip). Returns ``(columns, data_rows, truncated, masked_columns)`` on success, or the typed + ``Refusal`` on a guard refusal / execution failure — the SAME refusal the subprocess path produces + (``GuardRefused`` now carries the typed ``Refusal``, incl. the model-safety detail), so the two + paths build an identical refused Envelope. + + ``masked_columns`` (ACE-041) rides on the guarded ``ExecResult``: ``execute_guarded`` has already + redacted the sensitive values in ``result.rows``, so the rows textualized below are the redacted + ones; the refs travel through so the Envelope records ``applied[{mask}]``. Rows are textualized to match the subprocess CSV wire (``None`` → ``""``, else ``str``) so the two paths return observably identical JSON. Native-typed rows are a deliberately deferred decision @@ -996,10 +1024,11 @@ def _run_in_process( columns = list(result.columns) data_rows = [["" if v is None else str(v) for v in row] for row in result.rows] truncated = result.truncated + masked_columns = list(result.masked_columns) if max_rows is not None and len(data_rows) > max_rows: # backstop, matches the subprocess branch data_rows = data_rows[:max_rows] truncated = True - return columns, data_rows, truncated + return columns, data_rows, truncated, masked_columns def _refusal_from_stderr(stderr: str | None, returncode: int) -> Refusal: @@ -1125,10 +1154,11 @@ def tool_execute_sql(args: dict[str, Any]) -> str: sql=sql, execution_ms=execution_ms, ) - columns, data_rows, truncated = outcome + columns, data_rows, truncated, masked_columns = outcome return _finalize_execution( columns, data_rows, truncated, profile=profile, sql=sql, execution_ms=execution_ms, args=args, audit_id=audit_id, + masked_columns=masked_columns, ) # The model safety pass (fan/chasm pre-flight + default_filters) runs inside @@ -1180,10 +1210,14 @@ def tool_execute_sql(args: dict[str, Any]) -> str: if max_rows is not None and len(data_rows) > max_rows: data_rows = data_rows[:max_rows] truncated = True + # The redacted VALUES are already in `data_rows` (the executor masked result.rows before emitting + # the CSV); this reads the `{"masked": …}` marker so the Envelope records the applied[{mask}] note. + masked_columns = _executor_masked(proc.stderr) return _finalize_execution( columns, data_rows, truncated, profile=profile, sql=sql, execution_ms=execution_ms, args=args, audit_id=audit_id, + masked_columns=masked_columns, ) From 8ae2c79e854e4823b3ae7c17fbd652fed311101f Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 16:41:23 -0700 Subject: [PATCH 08/11] ACE-041 slice3: flip maskable-refuse tests to mask; widen _model_safety stubs; sync-lib - test_ah012: real-PII maskable projection now MASKS (query runs, value redacted, applied[{mask}] recorded); added an untraceable-projection refuse counterpart. - test_sample_database: raw-PII bare column now REDACTED end-to-end (token on stdout, {masked} marker on stderr); added an untraceable-still-refuses counterpart. - Mechanical: _model_safety callers/stubs widened 2->3 tuple (ace028, ace051, ah012). - sync-lib: vendored plugins/agami/lib/execute_sql.py byte-identical to source. These are spec-driven data-protection behaviour changes; NO F9 safety test modified. Co-Authored-By: Claude Opus 4.8 (1M context) --- plugins/agami/lib/execute_sql.py | 167 ++++++++++++++++++++---- tests/test_ace028_in_process_default.py | 2 +- tests/test_ace051_fail_closed.py | 18 +-- tests/test_ah012_executor_seam.py | 75 ++++++++--- tests/test_sample_database.py | 23 +++- 5 files changed, 230 insertions(+), 55 deletions(-) diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 341e2ca1..f48fbb2d 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -53,7 +53,7 @@ from typing import TYPE_CHECKING, Any import agami_paths -from guardrail import Refusal, Verdict +from guardrail import Refusal, Verdict, policy if TYPE_CHECKING: # ``Executor`` is the 5th port; imported only for type-checkers. At runtime ``execute_sql`` never @@ -113,6 +113,30 @@ class ExecResult: columns: list[str] rows: list[tuple] truncated: bool = False + # ACE-041: the DECLARED "table.column" refs of any OUTPUT columns redacted to REDACTION_TOKEN by + # the data-protection mask plan (empty for an unmasked result). The redacted VALUES already live + # in `rows`; this carries the `applied[{mask}]` note both surfaces record. Immutable default `()`. + masked_columns: tuple[str, ...] = () + + +# ACE-041: the string a masked sensitive value is replaced with in every row of a masked output +# column. A single named constant so the token is defined once for the redactor and the tests. +# ACE-041: token pending author confirm +REDACTION_TOKEN = "***" + + +@dataclass(frozen=True) +class MaskPlan: + """The post-execution redaction the data-protection gate proved safe: which OUTPUT column + positions to redact (`indices`, 0-based, aligned with the result's columns and every set-op arm) + and the DECLARED "table.column" refs those positions carry (`columns`, for the `applied[{mask}]` + note). Built by ``_model_safety`` only when ``guardrail.policy`` returns ``mask`` (every offending + projection is a deterministically-maskable 1:1 image of one output column), carried to + ``execute_guarded`` which applies it to ``result.rows``. Frozen — one plan per guarded call.""" + + indices: tuple[int, ...] + columns: tuple[str, ...] + token: str = REDACTION_TOKEN class ExecutorError(Exception): @@ -837,11 +861,21 @@ def _collect_cursor(cur: Any) -> ExecResult: return ExecResult(columns=columns, rows=[tuple(r) for r in fetched[:cap]], truncated=truncated) +def _flag_masked(columns: tuple[str, ...]) -> None: + """Signal which OUTPUT columns were PII-redacted to the subprocess caller — a non-error + ``{"masked": [...]}`` marker on stderr, the SAME metadata channel as ``{"truncated": …}`` (not a + new protocol). The redacted VALUES already ride the CSV on stdout; this line carries only the + ``applied[{mask}]`` note the parent records on the Envelope (ACE-041 slice 3).""" + sys.stderr.write(json.dumps({"masked": list(columns)}) + "\n") + + def _emit_result_csv(result: ExecResult) -> None: """Serialize an ``ExecResult`` to stdout as CSV — the subprocess/CLI wire. Byte-for-byte what the old inline cursor→CSV writer produced: header row then data rows, and a truncation marker on stderr when capped. This is the *single, final* text serialization for the fork path; the - in-process path skips it and returns the native rows straight to the tool edge.""" + in-process path skips it and returns the native rows straight to the tool edge. When the rows were + PII-masked, the redacted VALUES are already in ``result.rows`` (masked in ``execute_guarded`` + before this call), so the CSV carries them; a ``{"masked": …}`` stderr marker names the columns.""" if not result.columns: # cursor had no description → wrote nothing (e.g. a non-row statement) return writer = csv.writer(sys.stdout) @@ -850,6 +884,8 @@ def _emit_result_csv(result: ExecResult) -> None: writer.writerow(row) if result.truncated: _flag_truncated(_resolve_row_cap()) + if result.masked_columns: + _flag_masked(result.masked_columns) def _write_cursor_csv(cur: Any) -> None: @@ -920,15 +956,23 @@ def _refusal_from_verdict(kind: str, verdict: Verdict) -> Refusal: return Refusal(kind=kind, reason=verdict.detail, remediation=verdict.remediation) -def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusal | None]: - """Semantic-model safety pass before execution: fan-trap / chasm-trap pre-flight - + default_filters auto-application, over a model resolved from the DB (hosted) or disk (local). +def _model_safety( + sql: str, profile: str, area: str | None +) -> tuple[str, Refusal | None, MaskPlan | None]: + """Semantic-model safety pass before execution: fan-trap / chasm-trap pre-flight + scope + PII + + default_filters auto-application, over a model resolved from the DB (hosted) or disk (local). - Returns (sql_to_run, refusal). ``refusal`` is None to continue, or a typed ``Refusal`` the caller - (``execute_guarded``) raises as a ``GuardRefused`` — so both the subprocess and in-process paths - build the same envelope from it. Inert (returns the SQL unchanged) when the model package isn't - importable, or — on the LOCAL path only — when there is no model yet. On the HOSTED path a model - that can't be resolved fails closed (refuses), never runs unguarded (ACE-051).""" + Returns ``(sql_to_run, refusal, mask_plan)``: + - ``refusal`` is None to continue, or a typed ``Refusal`` the caller (``execute_guarded``) + raises as a ``GuardRefused`` — so both the subprocess and in-process paths build the same + envelope from it. + - ``mask_plan`` (ACE-041 slice 3) is None unless the PII gate proved a *maskable* sensitive + projection: the query then PROCEEDS and ``execute_guarded`` redacts the named output columns + post-execution. An *untraceable* sensitive projection sets ``refusal`` instead (fail-closed). + + Inert (SQL unchanged, no refusal, no plan) when the model package isn't importable, or — on the + LOCAL path only — when there is no model yet. On the HOSTED path a model that can't be resolved + fails closed (refuses), never runs unguarded (ACE-051).""" try: from semantic_model import runtime as RT except Exception: @@ -942,8 +986,8 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa "model_unavailable", "semantic-model package not importable; refusing to run unguarded on the " "hosted server", - ) - return sql, None # local: model package not available -> no-op + ), None + return sql, None, None # local: model package not available -> no-op org = _resolve_guard_model(profile) if org is None: @@ -954,8 +998,8 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa "model_unavailable", "no semantic model could be resolved (checked DB and disk); refusing to run " "unguarded on the hosted server", - ) - return sql, None # local: no model yet -> no-op (unchanged) + ), None + return sql, None, None # local: no model yet -> no-op (unchanged) # 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 @@ -968,41 +1012,79 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa # so the fan/chasm and sensitive checks below only evaluate in-scope tables. ts = RT.check_table_scope(sql, org, ctx=ctx) if ts is not None: - return sql, _refusal_from_verdict("table_out_of_scope", ts) + return sql, _refusal_from_verdict("table_out_of_scope", ts), 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, ctx=ctx) if star is not None: - return sql, _refusal_from_verdict("select_star", star) + return sql, _refusal_from_verdict("select_star", star), 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, ctx=ctx) if cs is not None: - return sql, _refusal_from_verdict("column_out_of_scope", cs) + return sql, _refusal_from_verdict("column_out_of_scope", cs), None pf = RT.pre_flight_check(sql, org, ctx=ctx) if pf.risk and pf.action == "refuse": - return sql, Refusal("preflight_refused", pf.reason, pf.suggestion or "") + return sql, Refusal("preflight_refused", pf.reason, pf.suggestion or ""), 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. + # Sensitive-column (PII) fail-closed on UNPARSEABLE SQL when the model declares PII (ACE-041 + # slice 3): a statement sqlglot cannot parse cannot be proven free of a raw sensitive projection. + # Unlike the other model guards (which degrade to allow — that fail-open is ACE-037's to close; + # ACE-037 is NOT on this branch), the PII path refuses a parse FAILURE outright rather than run a + # possibly-PII-projecting query. Only a parse failure with sqlglot PRESENT (ctx is not None and + # ctx.tree is None) is closed here; sqlglot-unavailable (ctx is None) stays the ACE-037-owned + # fail-open, and a model with no sensitive columns has nothing to protect. + if ctx is not None and ctx.tree is None and ctx.sensitive_by_table[1]: + return sql, Refusal( + "sensitive_columns", + "the query could not be parsed to verify it does not expose sensitive columns, so it " + "was refused.", + "Simplify the SQL so it parses, then retry.", + ), None + + # Sensitive-column (PII) guard — the deterministic chokepoint that protects PII identically for + # the agami-query skill, the local MCP server, and cron. A sensitive column may still be counted, + # filtered, grouped, or joined; only PROJECTING its raw value is offending. When it is, route the + # gate's verdict through the shared data-protection policy: a *maskable* projection (every + # offending projection is a clean 1:1 image of one output column) becomes a MASK plan and the + # query PROCEEDS with those columns redacted post-execution; an *untraceable* projection fails + # closed to a refusal (mirroring safety's "uncertainty ⇒ reject"). + mask_plan: MaskPlan | None = None sens = RT.check_sensitive_projection(sql, org, ctx=ctx) if sens.action == "refuse": - return sql, Refusal("sensitive_columns", sens.reason, sens.suggestion or "") + all_maskable = bool(sens.projections) and all( + p.maskable and p.output_index is not None for p in sens.projections + ) + verdict = Verdict( + cls="data_protection", + rule="sensitive_projection", + severity="high", + certainty="provable" if all_maskable else "uncertain", + detail=sens.reason, + remediation=sens.suggestion or "", + ) + # data_protection enforces in every tier (tier never downgrades it), so policy's default + # tier is correct here: provable → "mask", anything weaker → "reject" (fail closed). + if policy(verdict) == "mask": + indices = tuple(sorted({p.output_index for p in sens.projections})) + mask_plan = MaskPlan(indices=indices, columns=tuple(sens.columns)) + else: + return sql, Refusal("sensitive_columns", sens.reason, sens.suggestion or ""), None new_sql, applied = RT.apply_default_filters(sql, org, area=area, ctx=ctx) if applied: + # default_filters only add WHERE conditions, never change the projection list, so the mask + # plan's output indices stay valid against the rewritten SQL's result. sys.stderr.write(f"[agami] applied default_filters: {applied}\n") sql = new_sql - return sql, None + return sql, None, mask_plan # --------------------------------------------------------------------------- @@ -1143,18 +1225,49 @@ def execute_guarded( SQL both guards have passed. Raises ``GuardRefused`` carrying the typed ``Refusal`` on a refusal, and ``ExecutorError`` on a connect/run failure — so the subprocess ``main`` and the in-process MCP handler apply the same guard and build the same envelope. The row cap rides the request-scoped - ``_max_rows_override`` ContextVar the caller sets.""" + ``_max_rows_override`` ContextVar the caller sets. + + ACE-041: when the model-safety pass returns a ``MaskPlan`` (a *maskable* sensitive projection), + the query runs and the named output columns are redacted here, at the single shared post-execute + point — so masking covers every ``_run_`` AND any injected ``ports.Executor``, and composes + with the already-applied row cap (it redacts the capped rows in place).""" import sql_guard verdict = sql_guard.check_read_only(sql) if verdict is not None: raise GuardRefused(_refusal_from_verdict("permission", verdict), code=1) + mask_plan: MaskPlan | None = None if not no_safety: - sql, refusal = _model_safety(sql, profile, area) + sql, refusal, mask_plan = _model_safety(sql, profile, area) if refusal is not None: raise GuardRefused(refusal, code=1) creds = _load_credentials(profile) - return executor.execute(sql, creds, profile=profile) + result = executor.execute(sql, creds, profile=profile) + if mask_plan is not None: + result = _apply_mask_plan(result, mask_plan) + return result + + +def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: + """Redact the sensitive OUTPUT columns of an executor result post-execution — the SINGLE shared + masking point for EVERY executor (built-in ``_run_`` and an injected ``ports.Executor``). + Only the values at the plan's output indices become the token; every other value is untouched, + and the redacted rows plus the masked column refs are returned in a NEW ``ExecResult`` (it is + frozen). Runs on the already-row-capped rows, so it composes with the fetch bound rather than + fighting it. An index outside a row's width simply doesn't match (fail-safe — the value isn't + there to leak); by construction an output index equals its result-column position, so this holds.""" + if not plan.indices: + return result + mask_set = set(plan.indices) + redacted_rows = [ + tuple(plan.token if i in mask_set else v for i, v in enumerate(row)) for row in result.rows + ] + return ExecResult( + columns=result.columns, + rows=redacted_rows, + truncated=result.truncated, + masked_columns=plan.columns, + ) def main() -> int: diff --git a/tests/test_ace028_in_process_default.py b/tests/test_ace028_in_process_default.py index 12fdaf8c..15c947f9 100644 --- a/tests/test_ace028_in_process_default.py +++ b/tests/test_ace028_in_process_default.py @@ -105,7 +105,7 @@ def test_http_server_runs_in_process_by_default_no_fork(monkeypatch): monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None, None)) monkeypatch.setattr(tools.subprocess, "run", lambda *a, **k: pytest.fail("HTTP default must not fork")) out = json.loads(tools.tool_execute_sql({"sql": "SELECT 1 AS n", "datasource": "acme"})) diff --git a/tests/test_ace051_fail_closed.py b/tests/test_ace051_fail_closed.py index 6de552aa..dd84ea5a 100644 --- a/tests/test_ace051_fail_closed.py +++ b/tests/test_ace051_fail_closed.py @@ -76,7 +76,7 @@ def test_hosted_fail_closed_refuses_when_no_model(tmp_path, monkeypatch, capsys) monkeypatch.setenv("AGAMI_DB_URL", url) monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "no_artifacts")) - _, refusal = execute_sql._model_safety("SELECT id FROM orders", "acme", None) + _, refusal, _ = execute_sql._model_safety("SELECT id FROM orders", "acme", None) assert refusal is not None and refusal.kind == "model_unavailable" # refused (returned, not run) assert capsys.readouterr().err == "" # _model_safety returns the refusal, no stderr side-effect @@ -87,7 +87,7 @@ def test_local_missing_model_is_noop(tmp_path, monkeypatch): monkeypatch.delenv("APP_DATABASE_URL", raising=False) monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "empty")) - sql, code = execute_sql._model_safety("SELECT id FROM orders", "acme", None) + sql, code, _ = execute_sql._model_safety("SELECT id FROM orders", "acme", None) assert code is None and sql == "SELECT id FROM orders" # unchanged, guards inert @@ -98,11 +98,11 @@ def test_db_sourced_model_enforces_guards(tmp_path, monkeypatch, capsys): monkeypatch.setenv("AGAMI_DB_URL", url) monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "no_disk")) - _, refusal = execute_sql._model_safety("SELECT id FROM sqlite_master", "acme", None) + _, refusal, _ = execute_sql._model_safety("SELECT id FROM sqlite_master", "acme", None) assert refusal is not None and refusal.kind == "table_out_of_scope" # undeclared table, DB model assert capsys.readouterr().err == "" - sql, refusal = execute_sql._model_safety("SELECT id FROM orders", "acme", None) + sql, refusal, _ = execute_sql._model_safety("SELECT id FROM orders", "acme", None) assert refusal is None # a declared table with a named projection passes @@ -119,7 +119,7 @@ def verdict(hosted: bool, sql: str): else: monkeypatch.delenv("AGAMI_DB_URL", raising=False) monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "art")) # disk is the only source - _, code = execute_sql._model_safety(sql, "acme", None) + _, code, _ = execute_sql._model_safety(sql, "acme", None) capsys.readouterr() return code @@ -146,7 +146,7 @@ def test_hosted_falls_back_to_disk_when_db_has_no_model(tmp_path, monkeypatch, c monkeypatch.setenv("AGAMI_DB_URL", url) monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "art")) - _, refusal = execute_sql._model_safety("SELECT id FROM sqlite_master", "acme", None) + _, refusal, _ = execute_sql._model_safety("SELECT id FROM sqlite_master", "acme", None) assert refusal is not None and refusal.kind == "table_out_of_scope" # disk model, NOT model_unavailable assert capsys.readouterr().err == "" @@ -157,7 +157,7 @@ def test_hosted_db_load_error_falls_back_to_disk(tmp_path, monkeypatch): monkeypatch.setenv("AGAMI_DB_URL", "postgres://user:pw@127.0.0.1:1/nope") # unreachable monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "art")) - sql, code = execute_sql._model_safety("SELECT id FROM orders", "acme", None) + sql, code, _ = execute_sql._model_safety("SELECT id FROM orders", "acme", None) assert code is None # disk model resolved + guards passed the declared query @@ -168,7 +168,7 @@ def test_refusal_stderr_is_a_single_clean_json_object(tmp_path, monkeypatch, cap monkeypatch.setenv("AGAMI_DB_URL", "postgres://user:pw@127.0.0.1:1/nope") # unreachable → raises monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "no_disk")) # no disk model either - _, refusal = execute_sql._model_safety("SELECT id FROM orders", "acme", None) + _, refusal, _ = execute_sql._model_safety("SELECT id FROM orders", "acme", None) assert refusal is not None and refusal.kind == "model_unavailable" # The DB load error must not leak connection details into the refusal reason, and _model_safety # writes NOTHING to stderr — the JSON is emitted exactly once, by main()/execute_guarded. @@ -192,6 +192,6 @@ def boom(name, _globals=None, _locals=None, fromlist=(), level=0): monkeypatch.setattr(builtins, "__import__", boom) monkeypatch.setenv("AGAMI_DB_URL", "sqlite://" + str(tmp_path / "x.db")) - _, refusal = execute_sql._model_safety("SELECT id FROM orders", "acme", None) + _, refusal, _ = execute_sql._model_safety("SELECT id FROM orders", "acme", None) assert refusal is not None and refusal.kind == "model_unavailable" # fail closed, no DB load assert capsys.readouterr().err == "" diff --git a/tests/test_ah012_executor_seam.py b/tests/test_ah012_executor_seam.py index 197f18bf..03f2fb66 100644 --- a/tests/test_ah012_executor_seam.py +++ b/tests/test_ah012_executor_seam.py @@ -83,7 +83,7 @@ def test_model_safety_refusal_short_circuits_before_the_executor(monkeypatch): # A model-safety refusal short-circuits with its typed Refusal carried on the exception (so both # the subprocess and in-process paths build the same envelope); the executor must not run. monkeypatch.setattr( - execute_sql, "_model_safety", lambda s, p, a: (s, Refusal("preflight_refused", "fan-trap")) + execute_sql, "_model_safety", lambda s, p, a: (s, Refusal("preflight_refused", "fan-trap"), None) ) spy = _SpyExecutor() with pytest.raises(execute_sql.GuardRefused) as ei: @@ -99,7 +99,7 @@ def test_executor_receives_vetted_sql_and_resolved_creds(monkeypatch): # The default_filters rewrite happens in the model pass; the executor sees the POST-guard SQL and # the resolved datasource creds — never raw user input, never an unresolved profile. monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: ("SELECT 1 AS c /*vetted*/", None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: ("SELECT 1 AS c /*vetted*/", None, None)) spy = _SpyExecutor() result = execute_sql.execute_guarded("SELECT 1 AS c", "acme", "sales", executor=spy) @@ -218,7 +218,7 @@ def test_injected_executor_runs_in_process_with_vetted_sql_and_no_fork(monkeypat monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s + " /*vetted*/", None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s + " /*vetted*/", None, None)) # a fork here would be the REQ-002 violation the seam prevents — fail loudly if it happens. monkeypatch.setattr(tools.subprocess, "run", lambda *a, **k: pytest.fail("must not fork a subprocess")) @@ -249,7 +249,7 @@ def test_injected_executor_error_maps_to_the_same_envelope(monkeypatch): monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None, None)) class _Boom: def execute(self, vetted_sql, creds, *, profile): @@ -287,7 +287,7 @@ def _bad(profile): ) monkeypatch.setattr(execute_sql, "_load_credentials", _bad) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None, None)) tools.set_injected_executor(_SpyExecutor()) out = json.loads(tools.tool_execute_sql({"sql": "SELECT 1", "datasource": "acme"})) @@ -347,7 +347,7 @@ def test_injected_executor_model_safety_refusal_returns_the_typed_refusal(monkey monkeypatch.setattr( execute_sql, "_model_safety", - lambda s, p, a: (s, Refusal("table_out_of_scope", "table foo is not in the model")), + lambda s, p, a: (s, Refusal("table_out_of_scope", "table foo is not in the model"), None), ) fake = _SpyExecutor() tools.set_injected_executor(fake) @@ -359,11 +359,13 @@ def test_injected_executor_model_safety_refusal_returns_the_typed_refusal(monkey assert fake.calls == [] # refused before the executor -def test_injected_executor_real_pii_gate_produces_sensitive_columns_envelope(monkeypatch): - # A REAL sensitive-projection gate firing on a REAL model becomes a refused Envelope with kind - # "sensitive_columns" end-to-end (real gate -> _model_safety -> GuardRefused -> refused Envelope), - # closing the chain the synthetic-stderr test only stitched at the parse step. PII is top-severity, - # so pin the whole gate->refusal->envelope path, not just the stderr parser. The executor never runs. +def test_injected_executor_real_pii_gate_masks_a_maskable_projection(monkeypatch): + # BEHAVIOUR CHANGE (ACE-041 slice 3, was refuse before this slice): a REAL sensitive-projection + # gate firing on a REAL model for a *maskable* projection (a bare `email`, a clean 1:1 image of + # output column 0) now MASKS rather than refuses — the query RUNS and the value is redacted, with + # an `applied[{mask}]` note. This is the spec's data-protection behaviour, not a weakened safety + # gate: an UNTRACEABLE projection still refuses (see test_..._untraceable below and the ACE-041 + # masking suite). Pins the whole real gate -> policy -> mask -> redacted-envelope path end-to-end. pytest.importorskip("pydantic") pytest.importorskip("sqlglot") import tools @@ -390,13 +392,56 @@ def test_injected_executor_real_pii_gate_produces_sensitive_columns_envelope(mon ) monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: org) # real _model_safety runs - fake = _SpyExecutor() + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + fake = _SpyExecutor(result=execute_sql.ExecResult(columns=["email"], rows=[("alice@x.io",)], truncated=False)) tools.set_injected_executor(fake) out = json.loads(tools.tool_execute_sql({"sql": "SELECT email FROM users", "datasource": "acme"})) + assert out["status"] == "ok" # NOT refused — the query ran behind the guard + assert out["data"]["rows"] == [[execute_sql.REDACTION_TOKEN]] # the raw email was redacted, not leaked + assert "alice@x.io" not in json.dumps(out) # no raw sensitive value anywhere in the envelope + assert {"mask": "users.email"} in out["applied"] # the applied[{mask}] note is recorded + assert fake.calls # the executor RAN (masking is post-execution), then the value was redacted + + +def test_injected_executor_real_pii_gate_refuses_an_untraceable_projection(monkeypatch): + # The safety counterpart to the mask test: an UNTRACEABLE sensitive projection (the value buried + # in an expression) still fails closed to a refused Envelope with kind "sensitive_columns", and + # the executor never runs. Masking must NOT weaken this fail-closed refusal. + pytest.importorskip("pydantic") + pytest.importorskip("sqlglot") + import tools + from semantic_model import models as m + + org = m.Organization( + organization="o", + version=1, + subject_areas=[ + m.SubjectArea( + name="area", + description="d", + tables_defined=[ + m.Table( + name="users", schema="public", storage_connection="c", grain=["id"], + columns=[ + m.Column(name="id", type="integer"), + m.Column(name="email", type="string", sensitive=True), + ], + ) + ], + ) + ], + ) + monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: org) + fake = _SpyExecutor() + tools.set_injected_executor(fake) + + out = json.loads(tools.tool_execute_sql({"sql": "SELECT UPPER(email) FROM users", "datasource": "acme"})) + assert out["status"] == "refused" - assert out["refusal"]["kind"] == "sensitive_columns" # the REAL PII gate, not a synthetic line + assert out["refusal"]["kind"] == "sensitive_columns" # buried value → fail closed assert "data" not in out # a refusal carries no data (no raw sensitive values leak) assert fake.calls == [] # refused BEFORE the executor ran @@ -408,7 +453,7 @@ def test_injected_executor_textualizes_null_as_empty_at_the_tool_edge(monkeypatc monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None, None)) fake = _SpyExecutor(result=execute_sql.ExecResult(columns=["n", "s"], rows=[(1, None)], truncated=False)) tools.set_injected_executor(fake) @@ -422,7 +467,7 @@ def test_injected_executor_backstop_trims_to_max_rows(monkeypatch): monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) - monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None)) + monkeypatch.setattr(execute_sql, "_model_safety", lambda s, p, a: (s, None, None)) fake = _SpyExecutor(result=execute_sql.ExecResult(columns=["n"], rows=[(1,), (2,), (3,)], truncated=False)) tools.set_injected_executor(fake) diff --git a/tests/test_sample_database.py b/tests/test_sample_database.py index 8dd6b161..c6055bb6 100644 --- a/tests/test_sample_database.py +++ b/tests/test_sample_database.py @@ -280,10 +280,27 @@ def _run_execute_sql(art, sql): capture_output=True, text=True, env={**os.environ, "AGAMI_ARTIFACTS_DIR": str(art)}) -def test_execute_sql_refuses_raw_pii_both_paths(wired_artifacts): - """End-to-end through execute_sql.py — the path BOTH the skill and the MCP server - use — a raw PII projection is refused with a structured error, not leaked.""" +def test_execute_sql_masks_raw_pii_both_paths(wired_artifacts): + """End-to-end through execute_sql.py — the path BOTH the skill and the MCP server use. + + BEHAVIOUR CHANGE (ACE-041 slice 3, was `refuse` before this slice): a *maskable* raw-PII + projection (a bare `email` column) is now REDACTED, not refused — the query runs, the sensitive + VALUE is replaced with the token on stdout (so no raw email ever leaves the process), and a + `{"masked": …}` marker names the column on stderr. `full_name` (non-sensitive) is untouched. This + is the spec's data-protection masking, not a weakened safety gate (an untraceable projection + still refuses — see test_execute_sql_still_refuses_untraceable_pii).""" proc = _run_execute_sql(wired_artifacts, "SELECT full_name, email FROM customers LIMIT 5") + assert proc.returncode == 0, proc.stderr # the query RAN (masked, not blocked) + assert "@example.com" not in proc.stdout # no raw email leaked — the value was redacted + assert "***" in proc.stdout # the redaction token is present in the email column + assert '"masked"' in proc.stderr and "customers.email" in proc.stderr # the applied-mask marker + assert "Traceback" not in proc.stderr + + +def test_execute_sql_still_refuses_untraceable_pii(wired_artifacts): + """The safety counterpart: an UNTRACEABLE sensitive projection (the value buried in a function) + still fails closed — masking must not weaken the refuse path.""" + proc = _run_execute_sql(wired_artifacts, "SELECT UPPER(email) FROM customers LIMIT 5") assert proc.returncode == 1, proc.stderr assert "sensitive_columns" in proc.stderr assert "@example.com" not in proc.stdout # no raw email leaked From 70c764cdc142860515e25dcdfb7f5702a0aa753e Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 17:15:27 -0700 Subject: [PATCH 09/11] ACE-041 review: accurate applied[{mask}] note + known-limitation marker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finding 1 (correctness/fail-safe): _apply_mask_plan now builds the applied[{mask}] / masked_columns receipt via a new _mask_note helper. It keeps every declared sensitive table.column ref AND appends the output label (from result.columns) of any extra positionally-aligned position a set operation forced us to redact — e.g. "SELECT ssn, name FROM customers UNION ALL SELECT name, ssn FROM customers" redacts BOTH output columns (each carries a raw SSN in some arm) and the note now reads ("customers.ssn", "name") instead of under-reporting as ("customers.ssn",). A position is "already named" iff its output label case-folds to the bare column name of a declared ref; the normal single-column case is therefore unchanged and still reports exactly {"mask": "customers.ssn"} (success criterion 1). MaskPlan.columns is left as the declared refs so _model_safety's plan is unchanged; the enrichment happens once, post-execution, so both surfaces (in-process rows + subprocess {"masked": ...} stderr line) inherit it. Vendored copy synced. Finding 2 (known limitation, NOT fixed): documented the derived-table / CTE alias-rename bypass at semantic_model/runtime.py::_output_selects with an "# ACE-041 known limitation" comment pointing at the xfail regressions. Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/agami-core/src/execute_sql.py | 37 ++++++++++++++++++- .../agami-core/src/semantic_model/runtime.py | 12 +++++- plugins/agami/lib/execute_sql.py | 37 ++++++++++++++++++- 3 files changed, 81 insertions(+), 5 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index f48fbb2d..f3920a91 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -1248,6 +1248,35 @@ def execute_guarded( return result +def _mask_note( + declared: tuple[str, ...], indices: tuple[int, ...], out_columns: list[str] +) -> tuple[str, ...]: + """The ``applied[{mask}]`` receipt for a redaction — accurate about EVERY masked output position. + + ``declared`` are the model's sensitive "table.column" refs the gate named (``MaskPlan.columns``); + ``indices`` are the redacted 0-based output positions; ``out_columns`` are the executor's actual + output labels (``result.columns``), which for a set operation come from its FIRST arm. A single + sensitive projection redacts exactly one position whose own label IS that sensitive column, so the + note is just its declared ref (the spec's success criterion 1 — e.g. ``("customers.ssn",)``). + + A CROSS-POSITION set operation (e.g. ``SELECT ssn, name … UNION ALL SELECT name, ssn …``) redacts + a SECOND output position whose first-arm label is a non-sensitive name (``name``) — necessary, + because that merged output column carries a raw SSN in the second arm's rows. Naming that extra + position by its output label keeps the receipt honest: it reflects every redacted position, not + only the declared columns. A position is "already named by a declared sensitive column" iff its + output label case-folds to the bare column name of one of the ``declared`` refs; otherwise its + label is appended. This never DROPS a masked position (fail-safe for a receipt) and never renames + the single-column case (that position's label equals its declared column's bare name).""" + declared_bare = {d.rsplit(".", 1)[-1].lower() for d in declared} + note: list[str] = list(declared) + for i in sorted(indices): + if 0 <= i < len(out_columns): + label = out_columns[i] + if label.lower() not in declared_bare and label not in note: + note.append(label) + return tuple(note) + + def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: """Redact the sensitive OUTPUT columns of an executor result post-execution — the SINGLE shared masking point for EVERY executor (built-in ``_run_`` and an injected ``ports.Executor``). @@ -1255,7 +1284,11 @@ def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: and the redacted rows plus the masked column refs are returned in a NEW ``ExecResult`` (it is frozen). Runs on the already-row-capped rows, so it composes with the fetch bound rather than fighting it. An index outside a row's width simply doesn't match (fail-safe — the value isn't - there to leak); by construction an output index equals its result-column position, so this holds.""" + there to leak); by construction an output index equals its result-column position, so this holds. + + ``masked_columns`` is the ACCURATE receipt (``_mask_note``): the declared sensitive refs PLUS the + output label of any extra positionally-aligned position a set operation forced us to redact, so + the note reflects every masked output position, not only the declared columns (ACE-041 review).""" if not plan.indices: return result mask_set = set(plan.indices) @@ -1266,7 +1299,7 @@ def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: columns=result.columns, rows=redacted_rows, truncated=result.truncated, - masked_columns=plan.columns, + masked_columns=_mask_note(plan.columns, plan.indices, result.columns), ) diff --git a/packages/agami-core/src/semantic_model/runtime.py b/packages/agami-core/src/semantic_model/runtime.py index f80fe406..b9af8a35 100644 --- a/packages/agami-core/src/semantic_model/runtime.py +++ b/packages/agami-core/src/semantic_model/runtime.py @@ -531,7 +531,17 @@ def _output_selects(node: "exp.Expression") -> list["exp.Select"]: the table-scope fix). Nested subquery / CTE SELECTs are excluded on purpose: their projections feed an enclosing query, not the final result, so a sensitive column a WHERE-subquery projects but the outer query only filters on is not exposed and must - not be refused.""" + not be refused. + + # ACE-041 known limitation (pre-existing, NOT introduced by ACE-041): because only these + # OUTPUT-bearing SELECTs are walked and the sensitive match is name-based with no alias/lineage + # tracking, a sensitive column RENAMED inside a derived-table / CTE body escapes both mask and + # refuse — the outer query projects the renamed alias, which no longer name-matches the sensitive + # column, and this function never descends into the body that produced it. Two leaks are pinned as + # xfail regressions in tests/test_ace041_masking.py (`WITH t AS (SELECT ssn AS s FROM customers) + # SELECT s FROM t` and `SELECT z FROM (SELECT ssn AS z FROM customers) q`); they xfail today and + # flip to xpass when a future alias-lineage fix lands. Closing this needs real projection lineage + # through subquery/CTE bodies (a follow-up spec), which is deliberately out of ACE-041's scope.""" if isinstance(node, exp.Select): return [node] if isinstance(node, exp.SetOperation): # base of Union / Intersect / Except diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index f48fbb2d..f3920a91 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -1248,6 +1248,35 @@ def execute_guarded( return result +def _mask_note( + declared: tuple[str, ...], indices: tuple[int, ...], out_columns: list[str] +) -> tuple[str, ...]: + """The ``applied[{mask}]`` receipt for a redaction — accurate about EVERY masked output position. + + ``declared`` are the model's sensitive "table.column" refs the gate named (``MaskPlan.columns``); + ``indices`` are the redacted 0-based output positions; ``out_columns`` are the executor's actual + output labels (``result.columns``), which for a set operation come from its FIRST arm. A single + sensitive projection redacts exactly one position whose own label IS that sensitive column, so the + note is just its declared ref (the spec's success criterion 1 — e.g. ``("customers.ssn",)``). + + A CROSS-POSITION set operation (e.g. ``SELECT ssn, name … UNION ALL SELECT name, ssn …``) redacts + a SECOND output position whose first-arm label is a non-sensitive name (``name``) — necessary, + because that merged output column carries a raw SSN in the second arm's rows. Naming that extra + position by its output label keeps the receipt honest: it reflects every redacted position, not + only the declared columns. A position is "already named by a declared sensitive column" iff its + output label case-folds to the bare column name of one of the ``declared`` refs; otherwise its + label is appended. This never DROPS a masked position (fail-safe for a receipt) and never renames + the single-column case (that position's label equals its declared column's bare name).""" + declared_bare = {d.rsplit(".", 1)[-1].lower() for d in declared} + note: list[str] = list(declared) + for i in sorted(indices): + if 0 <= i < len(out_columns): + label = out_columns[i] + if label.lower() not in declared_bare and label not in note: + note.append(label) + return tuple(note) + + def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: """Redact the sensitive OUTPUT columns of an executor result post-execution — the SINGLE shared masking point for EVERY executor (built-in ``_run_`` and an injected ``ports.Executor``). @@ -1255,7 +1284,11 @@ def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: and the redacted rows plus the masked column refs are returned in a NEW ``ExecResult`` (it is frozen). Runs on the already-row-capped rows, so it composes with the fetch bound rather than fighting it. An index outside a row's width simply doesn't match (fail-safe — the value isn't - there to leak); by construction an output index equals its result-column position, so this holds.""" + there to leak); by construction an output index equals its result-column position, so this holds. + + ``masked_columns`` is the ACCURATE receipt (``_mask_note``): the declared sensitive refs PLUS the + output label of any extra positionally-aligned position a set operation forced us to redact, so + the note reflects every masked output position, not only the declared columns (ACE-041 review).""" if not plan.indices: return result mask_set = set(plan.indices) @@ -1266,7 +1299,7 @@ def _apply_mask_plan(result: ExecResult, plan: MaskPlan) -> ExecResult: columns=result.columns, rows=redacted_rows, truncated=result.truncated, - masked_columns=plan.columns, + masked_columns=_mask_note(plan.columns, plan.indices, result.columns), ) From b9cb3d675bdb448f1af0ca07fa9e0e9906c15a78 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 17:15:45 -0700 Subject: [PATCH 10/11] ACE-041 review: consolidate masking test corpus (findings 1-3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finding 1 — cross-position set operations (pin the behaviour + the enriched receipt): UNION ALL / UNION / INTERSECT / EXCEPT of (ssn, name) vs (name, ssn) redacts EVERY output position and leaks no raw SSN on both surfaces (in-process rows + CSV wire); the note is asserted as ("customers.ssn", "name"). The partial case "SELECT ssn, id FROM customers UNION ALL SELECT cust_id, id FROM orders" masks only column 0 and names only its declared ref. The single-column note is re-pinned unchanged, and _mask_note is unit-tested directly. Finding 2 — xfail the derived-table / CTE alias-rename bypass: "WITH t AS (SELECT ssn AS s FROM customers) SELECT s FROM t" and "SELECT z FROM (SELECT ssn AS z FROM customers) q" assert the DESIRED secure behaviour (outer projection masked OR refused). They xfail today (raw SSN flows through) with strict=False, so a future alias-lineage fix flips them to xpass as the tracking signal. Finding 3 — corpus gaps now asserted (were probe-only): a decimal sensitive column masks to "***" on both surfaces; a NULL at a masked index becomes the token deterministically (native + in-process); an empty result still records applied[{mask}] with zero rows; and a REAL subprocess fork ("python -m execute_sql" over an on-disk model + sqlite) asserts "***" in the emitted CSV, the raw value absent, and the {"masked": ...} stderr line present. The injected ports.Executor path is already covered by the existing slice-3 tool-edge test (not duplicated). Added an autouse fixture that resets the process-global injected executor between tests. No existing test weakened; F9 safety tests untouched. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_ace041_masking.py | 299 ++++++++++++++++++++++++++++++++++- 1 file changed, 297 insertions(+), 2 deletions(-) diff --git a/tests/test_ace041_masking.py b/tests/test_ace041_masking.py index 9ac5d28b..175295b9 100644 --- a/tests/test_ace041_masking.py +++ b/tests/test_ace041_masking.py @@ -17,6 +17,9 @@ from __future__ import annotations import json +import os +import sqlite3 +import subprocess import sys from pathlib import Path @@ -37,8 +40,9 @@ def _org() -> "m.Organization": - """`customers` has two sensitive columns (`ssn`, `email`); `orders` is non-sensitive for joins; - `x` carries a column literally named `ssn` that is NOT sensitive (the same-name decoy).""" + """`customers` has three sensitive columns (`ssn`, `email`, decimal `account_balance`); + `orders` is non-sensitive for joins and cross-datasource set-operation arms; `x` carries a column + literally named `ssn` that is NOT sensitive (the same-name decoy).""" customers = m.Table( name="customers", schema="public", storage_connection="c", grain=["id"], columns=[ @@ -46,6 +50,7 @@ def _org() -> "m.Organization": m.Column(name="name", type="string"), m.Column(name="ssn", type="string", sensitive=True), m.Column(name="email", type="string", sensitive=True), + m.Column(name="account_balance", type="decimal", sensitive=True), ], ) orders = m.Table( @@ -90,6 +95,17 @@ def _R(columns, rows, truncated=False): return execute_sql.ExecResult(columns=list(columns), rows=[tuple(r) for r in rows], truncated=truncated) +@pytest.fixture(autouse=True) +def _reset_injected_executor(): + """Isolate the process-global injected executor across tests (the tool-edge tests below register + a spy; leaving it set would pollute a later test). Additive — clears before and after each test.""" + import tools + + tools.set_injected_executor(None) + yield + tools.set_injected_executor(None) + + # --------------------------------------------------------------------------- # execute_guarded: a maskable projection PROCEEDS and the value is redacted. # --------------------------------------------------------------------------- @@ -333,3 +349,282 @@ def test_tool_execute_sql_untraceable_pii_is_refused_no_data(monkeypatch): assert out["refusal"]["kind"] == "sensitive_columns" assert "data" not in out assert spy.calls == [] + + +# --------------------------------------------------------------------------- +# Cross-position set operations (ACE-041 review, finding 1): a sensitive column that appears at a +# DIFFERENT output position in another arm forces that position to be redacted too — each merged +# output column carries a raw SSN in SOME arm's rows, so masking every such position is correct AND +# necessary. Two fixes are pinned here: (a) the masking behaviour (no raw SSN survives on either +# surface), and (b) the enriched `applied`/`masked_columns` note that names EVERY masked position. +# --------------------------------------------------------------------------- + + +# Each row carries a raw SSN in a DIFFERENT position, mirroring how a real DB merges the arms +# (arm 1 → (ssn, name); arm 2 → (name, ssn)). The output labels come from the first arm: (ssn, name). +_CROSS_POSITION_ROWS = [("111-22-3333", "Alice"), ("Bob", "999-88-7777")] + + +@pytest.mark.parametrize("op", ["UNION ALL", "UNION", "INTERSECT", "EXCEPT"]) +def test_cross_position_set_op_redacts_every_position_and_leaks_no_ssn(guarded, op): + # `SELECT ssn, name … SELECT name, ssn …`: output column 0 holds an SSN in arm 1's rows and + # column 1 holds an SSN in arm 2's rows, so BOTH output positions must be (and are) redacted. + spy = _SpyExecutor(_R(["ssn", "name"], _CROSS_POSITION_ROWS)) + result = guarded(f"SELECT ssn, name FROM customers {op} SELECT name, ssn FROM customers", spy) + tok = execute_sql.REDACTION_TOKEN + assert spy.calls # the query PROCEEDS (maskable), it is not refused + assert result.rows == [(tok, tok), (tok, tok)] # every position redacted + # No raw SSN survives at ANY output position, on the in-process (native rows) surface. + flat = [v for row in result.rows for v in row] + assert "111-22-3333" not in flat and "999-88-7777" not in flat + # Enriched note: the declared ref for column 0, plus column 1's output label (`name`) — the extra + # positionally-aligned position that carries raw SSN in arm 2 and is not named by a declared ref. + assert result.masked_columns == ("customers.ssn", "name") + + +def test_cross_position_union_redacts_no_ssn_on_the_csv_surface(guarded, capsys): + # The subprocess wire (CSV on stdout + `{"masked": …}` on stderr) must also carry no raw SSN and + # the enriched note. Run the guarded chain, then serialize the masked result the way the fork does. + spy = _SpyExecutor(_R(["ssn", "name"], _CROSS_POSITION_ROWS)) + result = guarded("SELECT ssn, name FROM customers UNION ALL SELECT name, ssn FROM customers", spy) + execute_sql._emit_result_csv(result) + captured = capsys.readouterr() + assert "***" in captured.out + assert "111-22-3333" not in captured.out and "999-88-7777" not in captured.out + masked = [json.loads(x)["masked"] for x in captured.err.splitlines() if '"masked"' in x] + assert masked == [["customers.ssn", "name"]] # the enriched note rides the stderr metadata channel + + +def test_cross_position_partial_only_names_the_masked_position(guarded): + # `SELECT ssn, id FROM customers UNION ALL SELECT cust_id, id FROM orders`: only output column 0 + # holds a sensitive value (ssn in arm 1; cust_id in arm 2 is NOT sensitive), so only column 0 is + # redacted. The whole column 0 is redacted (arm 2's non-sensitive cust_id is collateral, which is + # correct — no raw SSN can survive), and the note names just the one declared masked column. + spy = _SpyExecutor(_R(["ssn", "id"], [("111-22-3333", 1), (42, 2)])) + result = guarded("SELECT ssn, id FROM customers UNION ALL SELECT cust_id, id FROM orders", spy) + tok = execute_sql.REDACTION_TOKEN + assert result.rows == [(tok, 1), (tok, 2)] # column 0 redacted; column 1 (id) untouched + flat = [v for row in result.rows for v in row] + assert "111-22-3333" not in flat + assert result.masked_columns == ("customers.ssn",) # only the one masked position, its declared ref + + +def test_cross_position_enriched_note_reaches_the_tool_edge_applied(monkeypatch): + # The enriched note surfaces on the in-process tool surface's `applied[{mask}]` list too: both the + # declared ref AND the extra output-label position are recorded, so the receipt is accurate. + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + + spy = _SpyExecutor(_R(["ssn", "name"], _CROSS_POSITION_ROWS)) + tools.set_injected_executor(spy) + out = json.loads(tools.tool_execute_sql( + {"sql": "SELECT ssn, name FROM customers UNION ALL SELECT name, ssn FROM customers", + "datasource": "acme"} + )) + + assert out["status"] == "ok" + assert {"mask": "customers.ssn"} in out["applied"] + assert {"mask": "name"} in out["applied"] # the extra positionally-aligned position is named + + +def test_single_column_note_is_unchanged_by_the_enrichment(guarded): + # Regression guard for spec success criterion 1: the normal single-arm case still reports EXACTLY + # the declared ref — the enrichment must not add the output label when it already names the column. + spy = _SpyExecutor(_R(["ssn"], [("111-22-3333",)])) + result = guarded("SELECT ssn FROM customers", spy) + assert result.masked_columns == ("customers.ssn",) + + +def test_mask_note_helper_is_accurate_and_deterministic(): + # Unit-pin the note builder directly. Single declared column whose label matches -> just the ref. + assert execute_sql._mask_note(("customers.ssn",), (0,), ["ssn"]) == ("customers.ssn",) + # Aliased single column (label differs from the bare name) -> the label is added alongside the ref, + # which is fail-safe for a receipt (it never drops a masked position, only ever adds disclosure). + assert execute_sql._mask_note(("customers.ssn",), (0,), ["taxid"]) == ("customers.ssn", "taxid") + # Cross-position set op: column 0 named by the declared ref, column 1 named by its output label. + assert execute_sql._mask_note(("customers.ssn",), (0, 1), ["ssn", "name"]) == ("customers.ssn", "name") + # An index past the output width simply contributes nothing (fail-safe, no crash). + assert execute_sql._mask_note(("customers.ssn",), (0, 9), ["ssn"]) == ("customers.ssn",) + + +# --------------------------------------------------------------------------- +# KNOWN LIMITATION (ACE-041 review, finding 2): a sensitive column RENAMED inside a derived-table / +# CTE body escapes both mask and refuse, because `_output_selects` excludes subquery/CTE bodies and +# the sensitive match is name-based with no alias/lineage tracking. This is PRE-EXISTING (not +# introduced by ACE-041). These xfail tests assert the DESIRED secure behaviour (the outer projection +# is masked OR refused); they xfail today and will flip to xpass when a future alias-lineage fix +# lands — that flip is the tracking signal. Do NOT "fix" them by weakening the assertion. See the +# `# ACE-041 known limitation` comment in semantic_model/runtime.py::_output_selects. +# --------------------------------------------------------------------------- + +_ALIAS_LINEAGE_BYPASS = [ + "WITH t AS (SELECT ssn AS s FROM customers) SELECT s FROM t", + "SELECT z FROM (SELECT ssn AS z FROM customers) q", +] + + +@pytest.mark.xfail( + reason="ACE-041 known limitation: no alias/lineage tracking through derived-table/CTE bodies; " + "pre-existing in the sensitive-projection gate. Follow-up spec required.", + strict=False, +) +@pytest.mark.parametrize("sql", _ALIAS_LINEAGE_BYPASS) +def test_alias_rename_through_derived_body_should_be_masked_or_refused(guarded, sql): + # DESIRED secure behaviour: the outer projection of the renamed sensitive column is either masked + # or the whole query is refused. Today it is neither (the guard allows it and the raw value flows + # straight through), so this xfails; a future lineage fix flips it to xpass. + spy = _SpyExecutor(_R(["out"], [("111-22-3333",)])) + try: + result = guarded(sql, spy) + except execute_sql.GuardRefused: + return # refusing the query is a secure (desired) outcome + # It ran: the renamed sensitive value MUST be redacted and the raw value MUST NOT survive. + flat = [v for row in result.rows for v in row] + assert execute_sql.REDACTION_TOKEN in flat + assert "111-22-3333" not in flat + + +# --------------------------------------------------------------------------- +# Corpus gaps (ACE-041 review, finding 3): cases that previously passed only by probe — now asserted. +# --------------------------------------------------------------------------- + + +def test_numeric_sensitive_column_masks_on_both_surfaces(guarded, capsys): + # A numeric (typed) sensitive column projected bare is redacted to the token with no crash, on the + # in-process surface AND the CSV wire (the token replaces the value regardless of its Python type). + spy = _SpyExecutor(_R(["account_balance"], [(1234.56,), (9876,)])) + result = guarded("SELECT account_balance FROM customers", spy) + tok = execute_sql.REDACTION_TOKEN + assert result.rows == [(tok,), (tok,)] + assert result.masked_columns == ("customers.account_balance",) + # Subprocess/CSV surface: the token is serialized, never the raw number. + execute_sql._emit_result_csv(result) + out = capsys.readouterr().out + assert "***" in out + assert "1234.56" not in out and "9876" not in out + + +def test_null_value_at_a_masked_index_becomes_the_token(guarded): + # A NULL (None) at a masked index is redacted deterministically to the token, exactly like any + # other value — the redactor replaces by position, not by inspecting the value. + spy = _SpyExecutor(_R(["ssn"], [(None,), ("111-22-3333",)])) + result = guarded("SELECT ssn FROM customers", spy) + tok = execute_sql.REDACTION_TOKEN + assert result.rows == [(tok,), (tok,)] # the NULL row and the value row both become the token + assert result.masked_columns == ("customers.ssn",) + + +def test_null_at_masked_index_textualizes_to_the_token_in_process(monkeypatch): + # On the in-process tool edge, a masked NULL surfaces as the token (not "" — it was replaced before + # the None→"" textualization), so a redacted NULL is indistinguishable from a redacted value. + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + + spy = _SpyExecutor(_R(["ssn"], [(None,)])) + tools.set_injected_executor(spy) + out = json.loads(tools.tool_execute_sql({"sql": "SELECT ssn FROM customers", "datasource": "acme"})) + assert out["data"]["rows"] == [[execute_sql.REDACTION_TOKEN]] + + +def test_empty_result_with_maskable_projection_records_mask_and_no_rows(guarded): + # An empty result set through a maskable projection: the query proceeds, returns zero rows with no + # error, and the mask note is STILL recorded (asserting the actual behaviour — the receipt names + # the column that WOULD have been redacted even though no value was present). + spy = _SpyExecutor(_R(["ssn"], [])) + result = guarded("SELECT ssn FROM customers", spy) + assert spy.calls + assert result.rows == [] + assert result.masked_columns == ("customers.ssn",) + + +def test_empty_result_still_records_applied_mask_on_the_tool_edge(monkeypatch): + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: "acme") + monkeypatch.setattr(execute_sql, "_resolve_guard_model", lambda profile: _org()) + monkeypatch.setattr(execute_sql, "_load_credentials", lambda p: {"type": "sqlite", "path": ":memory:"}) + + spy = _SpyExecutor(_R(["ssn"], [])) + tools.set_injected_executor(spy) + out = json.loads(tools.tool_execute_sql({"sql": "SELECT ssn FROM customers", "datasource": "acme"})) + assert out["status"] == "ok" + assert out["data"]["rows"] == [] and out["data"]["row_count"] == 0 + assert {"mask": "customers.ssn"} in out["applied"] # the note is recorded even with zero rows + + +# --------------------------------------------------------------------------- +# A REAL subprocess fork end-to-end: `python -m execute_sql` over an on-disk model + a sqlite +# datasource, exercising the true fork path (main -> execute_guarded -> BUILTIN_EXECUTOR -> sqlite -> +# _emit_result_csv). Asserts `***` in the emitted CSV, the raw value absent, and the `{"masked": …}` +# stderr line present — the genuine subprocess wire, not a monkeypatched main(). +# --------------------------------------------------------------------------- + +def _write_disk_model(root: Path) -> None: + """Write a minimal v2 semantic-model tree declaring `customers.ssn` sensitive — the layout + `semantic_model.loader.load_organization` (and `execute_sql._resolve_guard_model`) read from + disk. Mirrors tests/test_semantic_model_loader.py::test_disk_round_trip, plus the sensitive flag.""" + import yaml + + (root / "datasources" / "c").mkdir(parents=True) + (root / "subject_areas" / "area" / "tables").mkdir(parents=True) + (root / "org.yaml").write_text(yaml.safe_dump({ + "organization": "AcmeCorp", "version": 1, + "storage_connections": [{"name": "c", "ref": "datasources/c/storage.yaml"}], + "subject_areas": ["subject_areas/area"], + })) + (root / "datasources" / "c" / "storage.yaml").write_text( + yaml.safe_dump({"name": "c", "storage_type": "PostgreSQL", "storage_config": {}})) + (root / "subject_areas" / "area" / "subject_area.yaml").write_text(yaml.safe_dump({ + "name": "area", "description": "d", + "tables": [{"storage_connection": "c", "schema": "public", "table": "customers"}], + })) + (root / "subject_areas" / "area" / "tables" / "customers.yaml").write_text(yaml.safe_dump({ + "name": "customers", "schema": "public", "storage_connection": "c", "grain": ["id"], + "description": "one line", + "columns": [{"name": "id", "type": "integer", "primary_key": True}, + {"name": "ssn", "type": "string", "sensitive": True}], + })) + (root / "subject_areas" / "area" / "relationships.yaml").write_text( + yaml.safe_dump({"relationships": []})) + + +def test_real_subprocess_fork_masks_ssn_in_the_emitted_csv(tmp_path): + pytest.importorskip("yaml") + # On-disk model (declares customers.ssn sensitive) + a real sqlite datasource with raw SSNs. + art = tmp_path / "art" + _write_disk_model(art / "acme") + db = tmp_path / "customers.db" + con = sqlite3.connect(db) + con.execute("CREATE TABLE customers (id INTEGER, ssn TEXT)") + con.executemany("INSERT INTO customers VALUES (?, ?)", [(1, "111-22-3333"), (2, "999-88-7777")]) + con.commit() + con.close() + + # Env-first credentials (a sqlite DSN) avoid the chmod-600 file gate; AGAMI_ARTIFACTS_DIR points the + # model resolver at the on-disk tree; the DB-url vars are cleared so `_hosted()` is False (disk model). + env = {**os.environ, + "AGAMI_ARTIFACTS_DIR": str(art), + "DATASOURCE_URL__ACME": f"sqlite:///{db}", + "PYTHONPATH": str(PKG_SRC) + os.pathsep + os.environ.get("PYTHONPATH", "")} + env.pop("AGAMI_DB_URL", None) + env.pop("APP_DATABASE_URL", None) + + proc = subprocess.run( + [sys.executable, "-m", "execute_sql", "--profile", "acme", "--sql", "SELECT ssn FROM customers"], + capture_output=True, text=True, env=env, timeout=60, + ) + + assert proc.returncode == 0, proc.stderr + # The emitted CSV carries the token, never the raw SSNs. + assert "***" in proc.stdout + assert "111-22-3333" not in proc.stdout and "999-88-7777" not in proc.stdout + # The `{"masked": …}` metadata line is present on stderr, naming the declared masked column. + masked = [json.loads(x)["masked"] for x in proc.stderr.splitlines() + if x.strip().startswith("{") and '"masked"' in x] + assert masked == [["customers.ssn"]] From 3d1b27fb7abc5353241715b4664f88e59220b0e0 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Mon, 13 Jul 2026 19:02:55 -0700 Subject: [PATCH 11/11] ACE-041: confirm redaction token '***'; wire ACE-062 follow-up ref Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/agami-core/src/execute_sql.py | 2 +- packages/agami-core/src/semantic_model/runtime.py | 2 +- plugins/agami/lib/execute_sql.py | 2 +- tests/test_ace041_masking.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index f3920a91..03733505 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -121,7 +121,7 @@ class ExecResult: # ACE-041: the string a masked sensitive value is replaced with in every row of a masked output # column. A single named constant so the token is defined once for the redactor and the tests. -# ACE-041: token pending author confirm +# ACE-041: token confirmed as "***" (author, 2026-07-13). REDACTION_TOKEN = "***" diff --git a/packages/agami-core/src/semantic_model/runtime.py b/packages/agami-core/src/semantic_model/runtime.py index b9af8a35..a400b433 100644 --- a/packages/agami-core/src/semantic_model/runtime.py +++ b/packages/agami-core/src/semantic_model/runtime.py @@ -541,7 +541,7 @@ def _output_selects(node: "exp.Expression") -> list["exp.Select"]: # xfail regressions in tests/test_ace041_masking.py (`WITH t AS (SELECT ssn AS s FROM customers) # SELECT s FROM t` and `SELECT z FROM (SELECT ssn AS z FROM customers) q`); they xfail today and # flip to xpass when a future alias-lineage fix lands. Closing this needs real projection lineage - # through subquery/CTE bodies (a follow-up spec), which is deliberately out of ACE-041's scope.""" + # through subquery/CTE bodies (follow-up spec ACE-062), which is deliberately out of ACE-041's scope.""" if isinstance(node, exp.Select): return [node] if isinstance(node, exp.SetOperation): # base of Union / Intersect / Except diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index f3920a91..03733505 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -121,7 +121,7 @@ class ExecResult: # ACE-041: the string a masked sensitive value is replaced with in every row of a masked output # column. A single named constant so the token is defined once for the redactor and the tests. -# ACE-041: token pending author confirm +# ACE-041: token confirmed as "***" (author, 2026-07-13). REDACTION_TOKEN = "***" diff --git a/tests/test_ace041_masking.py b/tests/test_ace041_masking.py index 175295b9..be9b6131 100644 --- a/tests/test_ace041_masking.py +++ b/tests/test_ace041_masking.py @@ -468,7 +468,7 @@ def test_mask_note_helper_is_accurate_and_deterministic(): @pytest.mark.xfail( reason="ACE-041 known limitation: no alias/lineage tracking through derived-table/CTE bodies; " - "pre-existing in the sensitive-projection gate. Follow-up spec required.", + "pre-existing in the sensitive-projection gate. Tracked as follow-up spec ACE-062.", strict=False, ) @pytest.mark.parametrize("sql", _ALIAS_LINEAGE_BYPASS)