From ce99bedcbff167eccd889fb892184fbf3525ee49 Mon Sep 17 00:00:00 2001 From: BADRI NARAYANAN Date: Sun, 6 Sep 2026 23:12:04 -0700 Subject: [PATCH] fix(chat): stop rejecting compound WHERE clauses; stop misleading the model into inventing a customer_id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two compounding bugs found by QA testing caused all 6 of the chat feature's starter questions to fail. sql_tool_validator.validate_query() walked every exp.Func node to check it against an allow-list, but in the installed sqlglot==30.18.0, AND/OR/XOR are modeled as Func subclasses (Connector -> Binary -> Func) -- so any compound WHERE clause was rejected as calling a disallowed function 'and'. Fixed by skipping exp.Connector before the allow-list check; find_all still recurses into a connector's operands, so `... AND pg_sleep(1)` is still rejected (covered by a new test -- the load-bearing property of this fix). Every other Func-typed operator (->, @>, ~, ^, TRY_CAST) stays subject to the allow-list, also covered. Separately, describe_schema() listed customer_id as a plain column on every curated view. Every view already scopes itself to the current customer via its own WHERE clause (ADR 19 §4.3), so the model never needs to filter on it -- but seeing it listed, it invented placeholder values ('your_customer_id') that failed as invalid UUID literals. Fixed structurally: customer_id is hidden from the schema description (still selectable in actual results -- harmless, it's always the caller's own id), with a one-line scope note in the tool output plus a matching system prompt rule as defense-in-depth. No migration; neither fix touches views, grants, or the allow-list's shape. New regression tests in test_sql_tool_validator.py and test_chat_security_perimeter.py, including a live compound-WHERE query against the real curated views. --- .../chat/chat_orchestration_service.py | 3 ++ .../services/chat/read_only_sql_executor.py | 18 ++++++- .../app/services/chat/sql_tool_validator.py | 9 ++++ .../test_chat_security_perimeter.py | 48 +++++++++++++++++++ backend/tests/unit/test_sql_tool_validator.py | 26 ++++++++++ 5 files changed, 103 insertions(+), 1 deletion(-) diff --git a/backend/app/services/chat/chat_orchestration_service.py b/backend/app/services/chat/chat_orchestration_service.py index d9a1737..11a343e 100644 --- a/backend/app/services/chat/chat_orchestration_service.py +++ b/backend/app/services/chat/chat_orchestration_service.py @@ -57,6 +57,9 @@ 6. When a tool result has more than one row, format it as a markdown table (a header row, then a `|---|---|` separator, then one row per record) rather than prose -- positions, transactions, and tax lots read far better as a table than as a paragraph of numbers. +7. Every view is already scoped to the customer you are answering for -- it returns their rows and + no one else's. Never reference `customer_id` in a `WHERE` clause and never invent an id value; + there is no id for you to supply. """ diff --git a/backend/app/services/chat/read_only_sql_executor.py b/backend/app/services/chat/read_only_sql_executor.py index a4acc77..00b0e36 100644 --- a/backend/app/services/chat/read_only_sql_executor.py +++ b/backend/app/services/chat/read_only_sql_executor.py @@ -20,6 +20,19 @@ log = get_logger(__name__) +# Every curated view bakes tenant scoping into its own WHERE clause (ADR 19 §4.3), so `customer_id` +# is the same value on every row the assistant can ever see. Advertising it as a plain column led +# the model to invent placeholder ids ('your_customer_id') that fail as UUID literals, so the +# schema description hides it. It stays selectable -- `SELECT *` returning the caller's own id is +# harmless. Safe only because chat is customer-only (S11 §2): `describe_schema`/`execute` below +# both pin `SessionRole.CUSTOMER`, so the views' adviser/admin escape hatch is unreachable here. +_HIDDEN_SCHEMA_COLUMNS: frozenset[str] = frozenset({"customer_id"}) + +_SCHEMA_SCOPE_NOTE = ( + "-- Every view below already returns only the signed-in customer's own rows. There is no " + "customer id for you to supply; never put customer_id in a WHERE clause." +) + def _stringify(value: object) -> str: return "" if value is None else str(value) @@ -47,10 +60,13 @@ def describe_schema(self, customer_id: uuid.UUID) -> str: columns_by_view: dict[str, list[str]] = {} for table_name, column_name in rows: + if column_name in _HIDDEN_SCHEMA_COLUMNS: + continue columns_by_view.setdefault(table_name, []).append(column_name) - return "\n".join( + described = "\n".join( f"{view}({', '.join(columns)})" for view, columns in sorted(columns_by_view.items()) ) + return f"{_SCHEMA_SCOPE_NOTE}\n{described}" def execute(self, sql: str, *, customer_id: uuid.UUID) -> SqlToolOutcome: """`execute_read_only_sql` (S11 §3). Validates first (ADR 19 §4.3).""" diff --git a/backend/app/services/chat/sql_tool_validator.py b/backend/app/services/chat/sql_tool_validator.py index d046c08..3cfaf2b 100644 --- a/backend/app/services/chat/sql_tool_validator.py +++ b/backend/app/services/chat/sql_tool_validator.py @@ -75,6 +75,15 @@ def validate_query(sql: str) -> ValidationResult: ) for func in statement.find_all(exp.Func): + # sqlglot models the boolean connectors AND/OR/XOR as `Func` subclasses (their MRO is + # And -> Connector -> Binary -> Func), even though they are operators, not calls -- so + # without this skip, any compound WHERE clause is rejected as calling a disallowed + # function 'and'. `find_all` still recurses into a connector's operands, so + # `... AND pg_sleep(1) IS NULL` is still caught and rejected; every other operator that + # also happens to be a `Func` (`->`, `@>`, `~`, `^`, TRY_CAST) stays subject to the + # allow-list below. + if isinstance(func, exp.Connector): + continue name = func.name.lower() if isinstance(func, exp.Anonymous) else func.sql_name().lower() if name not in _ALLOWED_FUNCTIONS: return ValidationResult( diff --git a/backend/tests/integration/test_chat_security_perimeter.py b/backend/tests/integration/test_chat_security_perimeter.py index c95f32c..b0599de 100644 --- a/backend/tests/integration/test_chat_security_perimeter.py +++ b/backend/tests/integration/test_chat_security_perimeter.py @@ -196,3 +196,51 @@ def test_curated_views_match_the_shared_allow_list( ).all() assert {row.table_name for row in rows} == CURATED_VIEW_NAMES + + +def test_describe_schema_never_advertises_customer_id( + owner_engine: Engine, chat_perimeter: None +) -> None: + """The views bake in tenant scoping (ADR 19 §4.3); advertising `customer_id` led the model to + invent placeholder ids that fail as UUID literals -- it is now hidden from the description.""" + customer_id = uuid.uuid4() + _insert_customer(owner_engine, customer_id) + + schema = ReadOnlySqlExecutor().describe_schema(customer_id) + # The scope note itself names `customer_id` deliberately (to warn the model off it) -- only + # the per-view column listing must never contain it. + columns_only = "\n".join(line for line in schema.splitlines() if "(" in line) + + assert "customer_id" not in columns_only + for view in CURATED_VIEW_NAMES: + assert f"{view}(" in schema + + +def test_read_only_sql_executor_runs_a_compound_where_starter_query( + owner_engine: Engine, chat_perimeter: None +) -> None: + """The month-to-date-return starter question, end to end: a compound WHERE now validates and + executes against the real curated views (regression test for the `exp.Connector` skip).""" + customer_id = uuid.uuid4() + _insert_customer(owner_engine, customer_id) + _insert_sub_period_return(owner_engine, customer_id) + + outcome = ReadOnlySqlExecutor().execute( + "SELECT sub_period_start, return_pct FROM v_period_return " + "WHERE sub_period_start >= DATE '2026-01-01' AND is_provisional = false", + customer_id=customer_id, + ) + + assert outcome.status == "success" + assert outcome.row_count == 1 + + +def test_validator_rejection_is_a_structured_outcome_not_an_exception( + owner_engine: Engine, chat_perimeter: None +) -> None: + """S11 §5.3: a rejection comes back as a tool result the agent can act on and retry within the + iteration cap, never as a raised error or a raw database message.""" + outcome = ReadOnlySqlExecutor().execute("SELECT * FROM posting", customer_id=uuid.uuid4()) + + assert outcome.status == "validator_rejected" + assert outcome.message diff --git a/backend/tests/unit/test_sql_tool_validator.py b/backend/tests/unit/test_sql_tool_validator.py index b329f87..a0154da 100644 --- a/backend/tests/unit/test_sql_tool_validator.py +++ b/backend/tests/unit/test_sql_tool_validator.py @@ -27,6 +27,14 @@ ("SELECT * FROM v_holdings /* trick */; SELECT 1", "comment-obfuscated multiple statements"), ("not even valid sql (((", "unparseable input"), ("", "empty input"), + ( + "SELECT * FROM v_holdings WHERE quantity_remaining > 0 AND pg_sleep(1) IS NULL", + "disallowed function nested inside an AND -- the Connector skip must not become a bypass", + ), + ("SELECT try_cast(symbol AS int) FROM v_holdings", "disallowed function try_cast"), + ("SELECT holdings_json -> 'x' FROM v_published_snapshot", "disallowed json-extract operator"), + ("SELECT value_end ^ 2 FROM v_period_return", "disallowed power operator"), + ("SELECT symbol FROM v_holdings WHERE symbol ~ 'A'", "disallowed regexp-match operator"), ] _VALID_QUERIES = [ @@ -41,6 +49,24 @@ "SELECT sale_date, sum(realized_gain_loss) FROM v_realized_gains " "GROUP BY sale_date ORDER BY sale_date" ), + ( + "SELECT symbol, quantity_remaining FROM v_holdings " + "WHERE quantity_remaining > 0 AND symbol = 'AAPL'" + ), + ( + "SELECT sub_period_start, return_pct FROM v_period_return " + "WHERE sub_period_start >= date_trunc('month', current_date) AND is_provisional = false" + ), + "SELECT * FROM v_dividends WHERE effective_date >= '2026-01-01' OR account_role = 'cash'", + ( + "SELECT symbol FROM v_tax_lots WHERE designation = 'long' " + "AND (quantity_remaining > 0 OR quantity_opened > 0)" + ), + ( + "SELECT symbol FROM v_holdings WHERE symbol LIKE 'A%' " + "AND acquired_at BETWEEN '2020-01-01' AND '2021-01-01'" + ), + "SELECT symbol FROM v_holdings WHERE NOT (symbol = 'A' AND quantity_remaining > 0)", ]