From 6a7e51b81bfcf04976eb66783ad4e4914dc05c88 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Tue, 6 Oct 2026 20:29:58 +0800 Subject: [PATCH] =?UTF-8?q?K3=20=E5=A4=8D=E6=A0=B8=E4=BF=AE=E6=AD=A3?= =?UTF-8?q?=EF=BC=88=E7=AC=AC=202=20=E8=BD=AE=EF=BC=89=EF=BC=9A*=20?= =?UTF-8?q?=E6=8A=8A=E9=81=AE=E7=BD=A9=E5=88=97=E5=B8=A6=E8=BF=9B=E5=88=AB?= =?UTF-8?q?=E7=9A=84=E5=88=97=E5=90=8D=E6=97=B6=E6=8C=89=E5=B1=95=E5=BC=80?= =?UTF-8?q?=E8=90=BD=E5=88=B0=E7=9A=84=E4=BD=8D=E7=BD=AE=E9=81=AE=EF=BC=9B?= =?UTF-8?q?=E5=87=BD=E6=95=B0=E5=8F=82=E6=95=B0=E9=87=8C=E7=9A=84=20JSON?= =?UTF-8?q?=20=E7=AE=AD=E5=A4=B4=E8=AE=A4=E4=BD=9C=E9=81=AE=E7=BD=A9?= =?UTF-8?q?=E5=88=97=EF=BC=9Bagent=20=E6=9F=A5=E5=88=B0=E7=9A=84=E9=81=AE?= =?UTF-8?q?=E7=BD=A9=E5=80=BC=E5=92=8C=E6=8A=A5=E5=91=8A=E9=87=8C=E7=9A=84?= =?UTF-8?q?=E5=88=86=E5=BC=80=E8=AE=B0=EF=BC=8C=E4=B8=8A=E9=99=90=E6=94=B6?= =?UTF-8?q?=E6=BB=A1=E4=B9=9F=E4=B8=8D=E4=B8=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - * 改名(复核点名的 blocker):WITH t(a, b, …) AS (SELECT * …) 这类列名表、UNION 后一支的 *、FROM 子查询或 LATERAL 的列名表、标量子查询和 ARRAY 子查询里的 *,以前 * 本身不算取自哪一列,原值原样交给核查 agent、落库,第二道防线也 认不出。现在 * 带出遮罩列的值时,它展开落到的那几列都算(masking._star_places):只展开一个来源、每一列都数得清的 按位置只遮取自遮罩列的那几列——子查询按选出项数,数据源里的表按查询当时冻结的表结构数(_schema_columns;工具边界 用数据源的 schema_cache,事后按快照的 schema_artifact,是同一份),几张列一样的月度表 SELECT * 连起来汇总,金额照常; 数不清的(没有表结构、列数对不上、几个来源)整段遮。表结构只用来数位置,不用来断定表里没有遮罩列。标量子查询里的 * 按 _value_columns 展开(展开数据源里的表一律算),EXISTS 里的 * 不算;递归的 CTE 不会绕不出来;带 EXCLUDE / REPLACE 的 * 不按位置数 - JSON 箭头:函数参数里没写表名的 col->'path' 被 sqlglot 解析成以 col 为参数的 lambda,_value_columns 把 lambda 的 参数名照列名收。SQLite json_each(profile->'$.tels')、Postgres jsonb_array_elements_text(profile->'tels')、MySQL JSON_UNQUOTE(profile->'$.tel') 这几种拆开、取出来的值在工具边界就遮 - 上限:run_verify_agent 以前先把调用方交进来的报告快照值并进 seen,收满 SECRETS_MAX 之后 agent 查到的值并不进去,后面 的查询、回答时的第二道防线都认不出。现在两份分开记(masking.SeenValues:given / found),查询工具按值遮、交回调用方 (VerifyOutcome.masked_values)时合起来,agent 查到的在前;反方审稿理由的按值遮也改成 agent 查到的先收 - 已知局限写进 inquiry 模块说明:表结构过期、列数没变却换了次序或改了名时可能数错;EXCEPT 后一支照 UNION 算(多遮) 测试:新增 test_masking_star_renames(真跑 SQLite 的各种 * 改名、按表结构数位置、表结构对不上、快照记着的表结构), 工具边界(交给模型的、落库的、通用工件接口)和审批质询端到端的 * 改名、JSON 箭头、上限收满;每处实现都做过去掉之后 测试变红的验证。 --- backend/app/engine/inquiry.py | 12 +- backend/app/engine/masking.py | 364 ++++++++++++++++-- backend/app/engine/nodes/challenge.py | 6 +- backend/app/engine/verify_agent.py | 55 +-- backend/app/tools/datasource.py | 15 +- backend/tests/test_challenge_node.py | 23 ++ backend/tests/test_inquiry_api.py | 77 ++++ backend/tests/test_masking_star_renames.py | 251 ++++++++++++ backend/tests/test_masking_table_functions.py | 52 +++ backend/tests/test_verify_agent_masking.py | 109 +++++- 10 files changed, 887 insertions(+), 77 deletions(-) create mode 100644 backend/tests/test_masking_star_renames.py diff --git a/backend/app/engine/inquiry.py b/backend/app/engine/inquiry.py index acb75d6c..1ca86ad7 100644 --- a/backend/app/engine/inquiry.py +++ b/backend/app/engine/inquiry.py @@ -56,7 +56,14 @@ IS NOT NULL THEN amount ELSE 0 END))同写在 WHERE 里,不算 - FROM、JOIN、LATERAL 里的表值函数(json_each、json_tree、unnest、jsonb_each_text、JSON_TABLE、LATERAL VALUES)拆开 遮罩列的,按出列名认它出的列(常见函数的出列名、写的列名表、别名、函数名,_function_columns),同名的别的列一并遮 - (多遮);不认得的函数又没写列名表的,整份快照都遮 + (多遮);不认得的函数又没写列名表的,整份快照都遮。函数参数里没写表名的 col->'path'(JSON 箭头,解析器当成以 col + 为参数的 lambda)按 col 认,真的 lambda(transform(arr, x -> …))的参数和遮罩列同名时也算(多遮) +- * 把遮罩列带进别的列名(WITH t(a, b, …) AS (SELECT * …) 这类列名表、UNION 后一支的 *、标量子查询里的 *):只展开一个 + 来源、那里每一列都数得清的,按位置只遮取自遮罩列的那几列——子查询按它的选出项数,数据源里的表按查询当时冻结的表结构 + 数(几张列一样的月度表 SELECT * 连起来汇总,金额照常)。表结构里查不到这张表、数出来的列数和实际展开的对不上,或者 + * 连着几个来源的,* 展开落到的那一段整段遮,同一段里的金额这类别的列一并遮(多遮);表结构过期、列数没变却换了次序 + 或改了名的,可能数错(同按表结构认标签列)。标量子查询里的 * 展开数据源里的表一律算取自遮罩列。EXCEPT 后一支只用来 + 去掉行,也照 UNION 的后一支算(多遮) - 按遮罩列分档的标签(CASE、IIF 的结果全是写死的常量):标签文字是 agent 写在 SQL 里的,不当遮罩值,SQL、回答正文里 照原样;标签列当分组键(每一档一行)照常显示,不分组的逐行标签列整列遮(格子里看不出某一笔落在哪一档),可 agent 在正文里写出「这一笔是大单」遮不住,按 id 分组再按档数个数也看得出(同按阈值查出某几行)。只有只装着标签的列不收: @@ -93,7 +100,8 @@ 报告快照里的遮罩值、agent 前面的查询里遮掉的值都看不到。只有一种照原样:先查了写着号码的备注(那时还不知道这个号码 是遮罩值),后来才查出号码——备注快照落库时原样,从通用接口打开照原样,质询工件接口照事件里记下的位置换。报告自己 的查询快照(写作时的查询)仍原样存着,从那里打开不经这里的遮罩 -- 遮罩值最多收 SECRETS_MAX 种写法,再多的按值认不出(按列遮照旧)。agent 自己查到的值先收,报告快照里的后收 +- 遮罩值最多收 SECRETS_MAX 种写法,再多的按值认不出(按列遮照旧)。agent 自己查到的值先收,报告快照里的后收:核查 + agent 把两份分开记(masking.SeenValues),查询时按值遮、交回调用方、回答时合起来,都是 agent 查到的在前 - 换了写法的号码只认数字之间夹着一到四个分隔符(空格、制表符、全角空格、不换行空格、窄空格这类空白,连字符、减号、 点、括号,含全角),带 +86 / 0086 / 86 前缀、全角数字:写成汉字(一三九……)、中间夹着别的字、拆成几处写的、跨行 写的认不出;和别的数字只隔着一个连字符或点连着的(13900000001-2、19914.74)不算号码;SQL 里不带引号的数不这样认 diff --git a/backend/app/engine/masking.py b/backend/app/engine/masking.py index 597d24fc..b5775c11 100644 --- a/backend/app/engine/masking.py +++ b/backend/app/engine/masking.py @@ -11,16 +11,19 @@ 审稿理由里的按值遮降为第二道防线。查询失败时报错原话里的值也换(mask_error_text) - 哪些列(_derived_hidden):从遮罩列本身起,按 SQL 一层层认出值取自遮罩列的结果列(别名、函数、子查询、WITH、UNION 的 - 另一支、和 * 写在同一支里没起别名的表达式、FROM 里拆开遮罩列的表值函数);COUNT 只数个数、窗口的分区键只决定次序、 - 按遮罩列分档又当分组键或套在汇总里数个数的、汇总里的 CASE 只在条件里用到的不算(_value_columns、_band_only)。认不出 - 落在哪一列的、SQL 解析不了又写着遮罩列的,整份快照的列都算。双盲复算按同一套认法在派盲算之前挡下取自遮罩列的指标 - (derived_columns) + 另一支、和 * 写在同一支里没起别名的表达式、FROM 里拆开遮罩列的表值函数、函数参数里的 JSON 箭头 col->'path');* 把 + 遮罩列带进别的列名的(写了列名表的派生来源、UNION 后一支的 *、标量子查询里的 *)按它展开落到的位置算(_star_places, + 展开数据源里的表的按查询当时冻结的表结构数位置,_schema_columns)。 + COUNT 只数个数、窗口的分区键只决定次序、按遮罩列分档又当分组键或套在汇总里数个数的、汇总里的 CASE 只在条件里用到的 + 不算(_value_columns、_band_only)。认不出落在哪一列的、SQL 解析不了又写着遮罩列的,整份快照的列都算。双盲复算按同一套 + 认法在派盲算之前挡下取自遮罩列的指标(derived_columns) - 只装着分档标签的列(_label_columns):标签是 agent 写在 SQL 里的字,不当遮罩值;* 展开数据源里的表时按冻结的表结构 认它带不带同名的列(_schema_tables) - 哪些值(_Secrets):遮罩列里出现过的值,按值认有门槛(SECRET_MIN、SECRET_DIGITS),换的时候先按词、按开头几个字筛出 正文里真有的那几个值,不现编上万个分支的正则;正文、文字格、SQL 的字符串和注释里换了写法的长号码(199-1474-2785、 +86 …、全角数字)按纯数字写法在集合里查(loose_spans);已经摆在读者眼前的值免遮,之后并进来的照样免(protected) - 快照遮好的样子(_masked_view),回答时才认得出、事后要照位置重放的地方(view_edits、replay_edits) +- 一次核查里认出的遮罩值(SeenValues):调用方交进来的和 agent 查询时遮掉的分开记,合起来时 agent 查到的在前 除 mask_query_result、换了写法的号码和这一轮新加的几样之外,这里的函数原来都在 engine/inquiry.py 里,名字不变, inquiry 照旧导出(别处按 inquiry 的名字引用)。 @@ -32,7 +35,7 @@ import functools import json import logging import re -from collections.abc import Iterable +from collections.abc import Callable, Iterable from dataclasses import dataclass, field from decimal import Decimal from typing import Any @@ -72,16 +75,26 @@ def _masked_word() -> str: # -------------------------------------------------------------------------- +#: _value_columns 里「值可能取自数据源里一张表的任何一列」:套在里面的子查询里的 * 展开了表,SQL 里看不出这张表有哪几列。 +#: 按它认(_reads)一律算取自遮罩列(宁可多遮);真实的列名不会是这个 +_ANY_COLUMN = "*" + + def _value_columns(node: Any) -> set[str]: """一个选出列的值取自哪些列(列名小写)。COUNT 里的只数个数,窗口的分区键、排序键只决定次序,结果全是写死常量的 分档(_band_only:当分组键,或者套在汇总里按条件数个数)只决定一行落在哪一档,都带不出值,不算;套在里面的子查询 只看它选出来的列(它的条件、连接里写的列带不出值)。汇总里的 CASE / IIF 只看各个结果:条件只决定这一行算不算进来, 同写在 WHERE 里(SUM(CASE WHEN buyer_phone IS NOT NULL THEN amount ELSE 0 END) 的值全部来自金额);不汇总的逐行 - CASE 照旧连条件一起算(每一行落在哪一支本身就是这一行的信息)。""" + CASE 照旧连条件一起算(每一行落在哪一支本身就是这一行的信息)。 + 套在里面的子查询选的是 *(标量子查询、ARRAY(SELECT * …)):值取自 * 展开的来源(_star_inputs),派生来源看它的各个 + 选出项,表值函数看参数,展开了数据源里的表的记成 _ANY_COLUMN;EXISTS 里的 * 只说有没有,不算。函数参数里的 + col->'path'(JSON 箭头)被解析器当成以 col 为参数的 lambda,参数名照列名收(真的 lambda 的参数多算一列无害)。""" from sqlglot import exp out: set[str] = set() stack = [node] + # 已经展开过的派生来源:递归的 CTE 里 * 展开它自己时不再绕回去 + expanded = {id(node)} while stack: item = stack.pop() if isinstance(item, exp.Column): @@ -94,13 +107,32 @@ def _value_columns(node: Any) -> set[str]: stack.extend(r for r in results if r is not None) elif isinstance(item, exp.Window): stack.append(item.this) + elif isinstance(item, exp.Lambda): + out.update(p.name.lower() for p in item.expressions if isinstance(p, exp.Identifier) and p.name) + stack.append(item.this) elif isinstance(item, exp.Select): - stack.extend(item.expressions) + for part in item.expressions: + if not _is_star(part): + stack.append(part) + elif not isinstance(item.parent, exp.Exists): + inputs = _star_inputs(part, item) + if inputs is None: + out.add(_ANY_COLUMN) + continue + fresh = [n for n in inputs if id(n) not in expanded] + expanded.update(id(n) for n in fresh) + stack.extend(fresh) else: stack.extend(item.iter_expressions()) return out +def _reads(names: set[str], columns: set[str]) -> bool: + """值取自 columns(_value_columns)的一项是不是取自 names:有一列对上,或者 * 展开了数据源里的表(_ANY_COLUMN,说不清 + 带出的是哪一列)。""" + return bool(names & columns) or _ANY_COLUMN in columns + + def _aggregated(node: Any) -> bool: """node 套在同一层 SELECT 的汇总函数里(SUM(…)、AVG(…) OVER (…) 也算)。""" from sqlglot import exp @@ -230,6 +262,26 @@ def _schema_tables(schema: Any) -> dict[str, set[str]] | None: return out +def _schema_columns(schema: Any) -> dict[str, list[str]] | None: + """冻结的表结构 → {表名(小写): 按表里次序的列名(小写)},名字的记法同 _schema_tables;几个库里同名、列又不一样的表 + 不记(说不清 * 展开的是哪一张)。* 展开数据源里的表时按它数位置(_star_columns);整份拿不到的为 None。""" + tables = schema.get("tables") if isinstance(schema, dict) else None + if not isinstance(tables, dict): + return None + out: dict[str, list[str]] = {} + clashed: set[str] = set() + for key, meta in tables.items(): + columns = meta.get("columns") if isinstance(meta, dict) else None + if not isinstance(columns, list) or not all(isinstance(c, dict) and c.get("name") is not None for c in columns): + continue + names = [str(c["name"]).lower() for c in columns] + for name in {str(key).lower(), str(key).lower().rsplit(".", 1)[-1], + str(meta.get("qualified") or key).lower().rsplit(".", 1)[-1]}: + if out.setdefault(name, names) != names: + clashed.add(name) + return {name: names for name, names in out.items() if name not in clashed} + + def _table_lacks(source: Any, name: str, memo: _LabelMemo) -> bool: """FROM 里的这个来源是数据源里的一张表,冻结的表结构里查得到它、又没有 name 这一列:* 展开它带不出 name。""" from sqlglot import exp @@ -577,6 +629,165 @@ def _places(branch: list[Any], width: int | None) -> list[int | None]: for i in range(len(branch))] +#: 按位置数 * 展开成的列最多往下数几层:再深的(递归的 CTE 绕回自己)按数不清算 +_STAR_DEPTH = 16 +#: * 上的改动(DuckDB 一类的 EXCLUDE / REPLACE / RENAME / ILIKE):展开成的列不再是来源那几列的原样,不按位置数 +_STAR_EDITS = ("except_", "replace", "rename", "ilike") + + +def _star_sources(star: Any, select: Any) -> list[Any] | None: + """select 里的这个 * 展开的是哪几个来源:没写表名的是 FROM、JOIN、LATERAL VIEW 里的全部,写了的(t.*)只有那一个; + 写的表名对不上任何来源(外层的表)为 None。""" + from sqlglot import exp + + sources = _from_sources(select) + [((lateral.alias_or_name or "").lower(), lateral) + for lateral in select.args.get("laterals") or []] + qualifier = (star.table.lower() or None) if isinstance(star, exp.Column) else None + if qualifier is None: + return [source for _, source in sources] + found = [source for alias, source in sources if alias == qualifier] + return found or None + + +def _star_inputs(star: Any, select: Any) -> list[Any] | None: + """* 带出的值取自哪里:要接着看的节点——派生来源的每一支(值取自它的各个选出项),表值函数、VALUES 这类来源本身 + (值取自参数)。展开了数据源里的表、写的表名对不上任何来源的为 None:SQL 里看不出这张表有哪几列,哪一列都可能带出来。""" + from sqlglot import exp + + sources = _star_sources(star, select) + if sources is None: + return None + out: list[Any] = [] + for source in sources: + derived = _derived_block(source) + if derived is not None: + out += _leaves(derived[0]) + elif isinstance(source, exp.Table) and isinstance(source.this, exp.Identifier): + return None + else: + out.append(source) + return out + + +def _star_feeds(star: Any, select: Any, names: set[str]) -> bool: + """select 里的这个 * 带出的值里有没有取自 names 的(展开了数据源里的表的一律算,见 _star_inputs)。""" + inputs = _star_inputs(star, select) + return inputs is None or any(_reads(names, _value_columns(node)) for node in inputs) + + +def _table_columns(source: Any, tables: dict[str, list[str]] | None) -> list[str] | None: + """FROM 里的这个来源是数据源里的一张表、冻结的表结构里查得到它时,它按次序的列名(_schema_columns);别的为 None。""" + from sqlglot import exp + + if not tables or not isinstance(source, exp.Table) or not isinstance(source.this, exp.Identifier): + return None + return tables.get(source.name.lower()) + + +def _single_source(star: Any, select: Any, depth: int) -> Any: + """* 只展开一个来源、展开成的就是那个来源的列原样(没带 EXCLUDE 这类改动)时,那个来源;别的为 None。""" + from sqlglot import exp + + node = star.this if isinstance(star, exp.Column) else star + if depth > _STAR_DEPTH or any(node.args.get(k) for k in _STAR_EDITS): + return None + sources = _star_sources(star, select) + return sources[0] if sources is not None and len(sources) == 1 else None + + +def _star_columns(star: Any, select: Any, names: set[str], memo: dict[int, Any], + tables: dict[str, list[str]] | None, depth: int = 0) -> list[bool] | None: + """* 展开成的每一列(按位置)取不取自 names:只展开一个来源时才数得清——派生来源按它的每一列(_block_columns),数据源 + 里的表按查询当时冻结的表结构(tables,_schema_columns)。表结构里查不到的表、几个来源(按同名列连起来的会合成一列)、 + 表值函数,或者 * 带着 EXCLUDE / REPLACE 这类改动的,为 None。""" + source = _single_source(star, select, depth) + if source is None: + return None + if (derived := _derived_block(source)) is not None: + return _block_columns(derived[0], names, memo, tables, depth + 1) + columns = _table_columns(source, tables) + return [column in names for column in columns] if columns is not None else None + + +def _block_columns(block: Any, names: set[str], memo: dict[int, Any], tables: dict[str, list[str]] | None, + depth: int = 0) -> list[bool] | None: + """一块查询(派生来源)的每一列(按位置)取不取自 names:每一支一项项数,* 按它展开成的几列数(_star_columns),任何 + 一支的那一列取自 names 就算;有一支数不清、几支不一样宽的为 None。memo 记下数过的(id(那块查询) → 结果)。""" + from sqlglot import exp + + if id(block) in memo: + return memo[id(block)] + memo[id(block)] = None + rows: list[list[bool]] = [] + for leaf in _leaves(block): + if not isinstance(leaf, exp.Select): + return None + row: list[bool] = [] + for item in leaf.expressions: + if not _is_star(item): + row.append(_reads(names, _value_columns(item))) + elif (spread := _star_columns(item, leaf, names, memo, tables, depth)) is not None: + row += spread + else: + return None + rows.append(row) + if not rows or len({len(row) for row in rows}) != 1: + return None + memo[id(block)] = [any(column) for column in zip(*rows)] + return memo[id(block)] + + +def _branch_names(select: Any, tables: dict[str, list[str]] | None, depth: int = 0) -> list[str | None] | None: + """一支(一条 SELECT)结果里按次序的列名:没起别名的表达式为 None(_output_name 认不准的);* 按它展开成的列名 + (只展开一个来源时:数据源里的表按表结构,派生来源按它写的列名表、第一支的列名)。有 * 说不清展开成哪几列的为 None。 + UNION 第一支是 * 时,后一支的值按位置落进的就是这几个名字(_source_names)。""" + from sqlglot import exp + + out: list[str | None] = [] + for item in select.expressions: + if not _is_star(item): + out.append(_output_name(item)) + continue + source = _single_source(item, select, depth) + if source is None: + return None + if (derived := _derived_block(source)) is not None: + block, listed = derived + first = _leaves(block)[0] + names = list(listed) if listed is not None else ( + _branch_names(first, tables, depth + 1) if isinstance(first, exp.Select) else None) + else: + names = _table_columns(source, tables) + if names is None: + return None + out += names + return out + + +def _star_places(select: Any, width: int | None, names: set[str], tables: dict[str, list[str]] | None = None + ) -> set[int] | None: + """select 这一支里的 * 把取自 names 的值带到了结果(width 列)的哪几列。* 展开的列落在第一个 * 起、最后一个 * 后面 + 那几项之前的那一段:只有一个 *、它展开成的每一列都数得清的(_star_columns:派生来源,表结构里查得到的表)只算取自 + names 的那几列,别的整段都算。数出来的列数和这一段对不上的(表结构过期了)也整段算。没有 * 带出这样的值为空;不知道 + 结果有几列、* 前后的几项就比结果还宽的为 None。 + 用在 * 展开的列换了名字的地方:写了列名表的派生来源(WITH t(a, b) AS (SELECT * FROM sales))、UNION 后一支的 * + (结果的列名取第一支的)。""" + branch = list(select.expressions) + stars = [i for i, item in enumerate(branch) if _is_star(item)] + # 展开数据源里的表的一律算带出了遮罩列的值:表结构只用来数位置,不用来断定这张表没有遮罩列(表结构可能过期) + if not any(_star_feeds(branch[i], select, names) for i in stars): + return set() + if width is None: + return None + first, end = stars[0], width - (len(branch) - 1 - stars[-1]) + if end < first: + return None + spread = _star_columns(branch[first], select, names, {}, tables) if len(stars) == 1 else None + if spread is not None and len(spread) == end - first: + return {first + k for k, hit in enumerate(spread) if hit} + return set(range(first, end)) + + def _blocks(root: Any) -> list[tuple[Any, str, list[str] | None]]: """SQL 里的每一块查询(一条 SELECT,或者 UNION 这类连起来的一串)和它的结果流到哪: - "result":最外层,结果就是快照的列; @@ -608,10 +819,16 @@ def _blocks(root: Any) -> list[tuple[Any, str, list[str] | None]]: return out -def _source_names(top: Any, listed: list[str] | None, names: set[str]) -> tuple[set[str], bool]: +def _source_names(top: Any, listed: list[str] | None, names: set[str], + tables: dict[str, list[str]] | None = None) -> tuple[set[str], bool]: """派生来源的结果里值取自遮罩列的那几列的名字,以及有没有认不出名字的(没起别名的表达式、两个 * 之间的项)。 - 写了列名表的按位置对上列名表;没写的,单独一条 SELECT 取各项自己的列名,UNION 这类按位置对上第一支的列名。""" - branches = _branches(top) + 写了列名表的按位置对上列名表;没写的,单独一条 SELECT 取各项自己的列名,UNION 这类按位置对上第一支的列名。 + * 带出遮罩列的值、展开的列又换了名字的(写了列名表、UNION 后一支的 *),它展开落到的那几列照位置对上(_star_places); + 单独一条 SELECT、第一支的 * 不改名,遮罩列照原名传出去,按名字就认得出。第一支的 * 展开成哪几列数得清的(_branch_names), + 后一支落进去的位置按它对上列名。tables 是查询当时冻结的表结构(_schema_columns),* 展开数据源里的表时按它数位置。""" + from sqlglot import exp + + selects = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)] found: set[str] = set() unknown = False @@ -622,28 +839,39 @@ def _source_names(top: Any, listed: list[str] | None, names: set[str]) -> tuple[ else: found.add(name) + def taken(select: Any, width: int | None, head: Callable[[int | None], str | None]) -> None: + branch = list(select.expressions) + for item, place in zip(branch, _places(branch, width)): + if not _is_star(item) and _reads(names, _value_columns(item)): + take(head(place)) + spread = _star_places(select, width, names, tables) + for place in sorted(spread) if spread is not None else [None]: + take(head(place)) + if listed is not None: - for branch in branches: - for item, place in zip(branch, _places(branch, len(listed))): - if names & _value_columns(item): - take(listed[place] if place is not None and place < len(listed) else None) + for select in selects: + taken(select, len(listed), lambda p: listed[p] if p is not None and p < len(listed) else None) return found, unknown - first = branches[0] if branches else [] + first = list(selects[0].expressions) if selects else [] for item in first: - if names & _value_columns(item): + if not _is_star(item) and _reads(names, _value_columns(item)): take(_output_name(item)) - if len(branches) > 1: - # 每一支一样宽:拿一支没有 * 的定宽度,* 后面的项才数得出位置 - width = next((len(b) for b in branches if not any(_is_star(i) for i in b)), None) - heads = {place: _output_name(item) for item, place in zip(first, _places(first, width)) if place is not None} - for branch in branches[1:]: - for item, place in zip(branch, _places(branch, width)): - if names & _value_columns(item): - take(heads.get(place) if place is not None else None) + if len(selects) > 1: + header = _branch_names(selects[0], tables) + if header is not None: + width, heads = len(header), dict(enumerate(header)) + else: + # 每一支一样宽:拿一支没有 * 的定宽度,* 后面的项才数得出位置 + width = next((len(s.expressions) for s in selects if not any(_is_star(i) for i in s.expressions)), None) + heads = {place: _output_name(item) for item, place in zip(first, _places(first, width)) + if place is not None} + for select in selects[1:]: + taken(select, width, lambda p: heads.get(p) if p is not None else None) return found, unknown -def _tainted(sql: str, kind: str | None, masked: set[str], width: int) -> tuple[set[str], set[int]] | None: +def _tainted(sql: str, kind: str | None, masked: set[str], width: int, tables: dict[str, list[str]] | None = None + ) -> tuple[set[str], set[int]] | None: """SQL 里指代遮罩列的名字(小写),和结果(width 列)里值取自遮罩列的那几列的位置。解析不了返回 None。 名字从遮罩列本身起,一层层传下去:选出列的值取自这些名字的(buyer_phone AS 手机号、UPPER(buyer_phone) AS x),它的 @@ -653,7 +881,10 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int) -> tuple[ 后面的从末尾倒着数。认不出的——两个 * 之间的项、派生来源里没起别名的表达式(外层的 * 会把它原样带出来,列名 又认不准)——整份快照的列都算(同解析不了)。 FROM、JOIN、LATERAL 里的表值函数(json_each、unnest、JSON_TABLE、LATERAL VALUES)参数里写着这些名字的,它出的列 - 都算(_function_columns):遮罩列是 JSON 时,json_each 拆开是最自然的写法。出哪几列说不清的整份快照都算。""" + 都算(_function_columns):遮罩列是 JSON 时,json_each 拆开是最自然的写法。出哪几列说不清的整份快照都算。 + * 展开的列换了名字(写了列名表的派生来源、UNION 后一支的 *):* 带出这些名字的值时,它展开落到的那几列照位置算 + (_star_places,展开数据源里的表的按 tables 数位置,数不清的整段算);标量子查询里的 * 按 _value_columns 算进外层的 + 表达式。""" from sqlglot import exp from app.data.sql_semantics import _parse, _read @@ -676,29 +907,37 @@ def _tainted(sql: str, kind: str | None, masked: set[str], width: int) -> tuple[ for select_ in selects: for item in select_.expressions: name = item.alias_or_name.lower() - if name and name not in names and names & _value_columns(item): + if name and not _is_star(item) and name not in names and _reads(names, _value_columns(item)): names.add(name) for top, listed in sources: - names |= _source_names(top, listed, names)[0] + names |= _source_names(top, listed, names, tables)[0] for source in functions: - if names & _value_columns(source): + if _reads(names, _value_columns(source)): columns = _function_columns(source) unreadable = unreadable or columns is None names |= columns or set() if len(names) == before: break - everything = unreadable or any(_source_names(top, listed, names)[1] for top, listed in sources) + everything = unreadable or any(_source_names(top, listed, names, tables)[1] for top, listed in sources) positions: set[int] = set() for top, role, _listed in blocks: if role != "result": continue - for branch in _branches(top): + branches = [leaf for leaf in _leaves(top) if isinstance(leaf, exp.Select)] + for k, select_ in enumerate(branches): + branch = list(select_.expressions) for item, place in zip(branch, _places(branch, width)): - if names & _value_columns(item): + if not _is_star(item) and _reads(names, _value_columns(item)): if place is None or place >= width: everything = True else: positions.add(place) + # UNION 后一支的 * 按位置落进第一支起的列名里;第一支、单独一条 SELECT 的 * 照原名带出,按名字认得出 + spread = _star_places(select_, width, names, tables) if k else set() + if spread is None or any(place >= width for place in spread): + everything = True + else: + positions |= spread return names, set(range(width)) if everything else positions except Exception: # noqa: BLE001 - 分析本身的缺陷不能让回答出错:按解析不了处理(宁可多遮) logger.warning("质询的 SQL 遮罩分析失败", exc_info=True) @@ -709,8 +948,10 @@ def _mentions(sql: str, names: Iterable[str]) -> bool: return any(re.search(rf"(? tuple[set[str], set[str]]: - """(这份快照要遮的列, SQL 里指代遮罩列的名字),都是小写。 +def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None, + tables: dict[str, list[str]] | None = None) -> tuple[set[str], set[str]]: + """(这份快照要遮的列, SQL 里指代遮罩列的名字),都是小写。tables 是查询当时冻结的表结构(_schema_columns),* 展开 + 数据源里的表时按它数位置;没给的按快照记着的表结构快照(schema_artifact)认,和工具边界遮的时候是同一份。 证据面板只按列名遮(hidden);agent 的查询起了别名、套了函数,结果列名就对不上。本仓库的查询惯例是起中文别名 (region AS 区域),所以按 SQL 认出改了名字的遮罩列(_tainted)。SQL 解析不了、里面又写着遮罩列的,整份快照的 @@ -720,7 +961,7 @@ def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None if not base or not isinstance(sql, str) or not sql.strip() or snapshot.get("source") == MERGE_SOURCE: return base, base columns = [str(c) for c in snapshot.get("columns") or []] - found = _tainted(sql, kind, base, len(columns)) + found = _tainted(sql, kind, base, len(columns), _snapshot_columns(snapshot) if tables is None else tables) if found is None: return (base | {c.lower() for c in columns}, base) if _mentions(sql, base) else (base, base) names, positions = found @@ -729,12 +970,13 @@ def _derived_hidden(snapshot: dict[str, Any], hidden: set[str], kind: str | None def derived_columns(snapshot: dict[str, Any], hidden: set[str], kind: str | None) -> set[str] | None: """按 SQL 认出值取自遮罩列(hidden,小写)的结果列(小写,列名直接对上的也算)。SQL 解析不了、合并查询的快照(不看 - SQL)为 None:认不准,宽严由调用方定(_derived_hidden 整份遮;双盲复算照常派盲算,由盲算自己的查询去认)。""" + SQL)为 None:认不准,宽严由调用方定(_derived_hidden 整份遮;双盲复算照常派盲算,由盲算自己的查询去认)。表结构同 + _derived_hidden,按快照记着的表结构快照认。""" sql = snapshot.get("sql") if not hidden or not isinstance(sql, str) or not sql.strip() or snapshot.get("source") == MERGE_SOURCE: return None columns = [str(c) for c in snapshot.get("columns") or []] - found = _tainted(sql, kind, set(hidden), len(columns)) + found = _tainted(sql, kind, set(hidden), len(columns), _snapshot_columns(snapshot)) if found is None: return None names, positions = found @@ -753,6 +995,23 @@ def _snapshot_labels(snapshot: dict[str, Any], kind: str | None, tables: dict[st return _label_columns(sql, kind, [str(c) for c in snapshot.get("columns") or []], tables) +def _snapshot_columns(snapshot: dict[str, Any]) -> dict[str, list[str]] | None: + """快照记着的表结构快照(schema_artifact,查询当时冻结的,和工具边界遮的时候用的是同一份)里按次序的列名。""" + artifact = snapshot.get("schema_artifact") + return _frozen_columns(artifact) if isinstance(artifact, str) else None + + +@functools.lru_cache(maxsize=64) +def _frozen_columns(artifact: str) -> dict[str, list[str]] | None: + """表结构快照工件里每张表按次序的列名(_schema_columns)。同 _frozen_tables:按内容寻址,记下来;取不到的为 None。""" + from app.core import artifact_store + + try: + return _schema_columns(artifact_store.load(artifact)) + except Exception: # noqa: BLE001 - 取不到就当不知道表结构(* 展开表的整段遮,宁可多遮) + return None + + @functools.lru_cache(maxsize=64) def _frozen_tables(artifact: str) -> dict[str, set[str]] | None: """表结构快照工件里每张表的列名(_schema_tables)。工件按内容寻址、不会变,记下来;取不到、被改过的为 None。""" @@ -1489,16 +1748,39 @@ class QueryMask: secrets: _Secrets +@dataclass +class SeenValues: + """一次核查里认出的遮罩值,只在内存里:given 是调用方交进来的(审批质询、反方审稿收的报告快照里的值),found 是 + agent 自己的查询在工具边界遮掉的(tools/datasource 每次查询并进来,add)。两份分开记:合起来(pooled)时 agent 查到的 + 在前,到了上限 SECRETS_MAX 挤掉的是报告里的值——agent 查到的就在它眼前这几份结果里,后面的查询、它写的回答和理由 + 最可能再碰到。以前先并报告的值,收满了上限,agent 查到的一个也并不进去。""" + + given: _Secrets = field(default_factory=_Secrets) + found: _Secrets = field(default_factory=_Secrets) + + def add(self, secrets: _Secrets) -> None: + """并进一次查询时遮掉的值(mask_query_result 交回的 secrets)。""" + self.found.merge(secrets) + + def pooled(self) -> _Secrets: + """合起来的另一份(改它不影响这两份):agent 查到的在前,调用方交进来的在后;调用方免遮的文字(问题、被审的那句 + 结论原文里写着的值)对两边都算。上万个值合一遍要几十毫秒,调用方放进线程。""" + out = self.found.copy() + out.merge(self.given) + return out + + def mask_query_result(payload: dict[str, Any], masked: Iterable[str], kind: str | None, known: _Secrets | None = None, *, schema: Any = None) -> QueryMask: """一份查询结果(tools/datasource 的查询工具交回的样子:columns / rows / column_types / sql)在交给核查 agent、落进 快照之前先遮。masked 是数据源设了遮罩的列,kind 是数据源的类型(SQL 按它的方言解析)。known 是此前已经认出的遮罩值: 这个 agent 前面的查询遮掉的,加上调用方交进来的(审批质询、反方审稿收的报告快照里的值)。schema 是数据源的表结构 - (schema_cache,和随快照冻结的是同一份),认分档标签时用(_label_columns 的 tables)。 + (schema_cache,和随快照冻结的是同一份):认分档标签(_label_columns 的 tables)、数 * 展开数据源里的表落在第几列 + (_derived_hidden 的 tables)时用。 - 整列遮:列名直接对上遮罩列的,加上按 SQL 认出值取自遮罩列的(_derived_hidden:别名、函数、子查询、WITH、 - SELECT *、表值函数;COUNT、只在条件里用到、分档当分组键的不算);SQL 解析不了、又写着遮罩列的,整份结果的列都遮。 - 这几列的格子换成「已遮罩」,类型记成文本 + SELECT *(展开的列换了名字的也算)、表值函数;COUNT、只在条件里用到、分档当分组键的不算);SQL 解析不了、又写着 + 遮罩列的,整份结果的列都遮。这几列的格子换成「已遮罩」,类型记成文本 - 别的列按值换:整格就是遮罩值的换掉,文字格里夹着的那一截换掉(换了写法的号码也认,_Secrets.loose_spans)。遮罩值 是这份结果里的加上 known:先查号码、再查写着号码的备注、没设遮罩的联系方式表(别的库里的也一样)时,后一份同样 遮好。门槛同按值遮(SECRET_DIGITS、SECRET_MIN);只装着分档标签的列里的标签不算遮罩值,这份结果自己写的标签也不 @@ -1514,7 +1796,7 @@ def mask_query_result(payload: dict[str, Any], masked: Iterable[str], kind: str if not (base or given) or not columns or not _query_shaped(payload): return QueryMask(payload, [], _Secrets()) try: - hidden = _derived_hidden(payload, base, kind)[0] if base else set() + hidden = _derived_hidden(payload, base, kind, _schema_columns(schema))[0] if base else set() own = _Secrets() labels = _snapshot_labels(payload, kind, _schema_tables(schema)) if any(c.lower() in hidden for c in columns): diff --git a/backend/app/engine/nodes/challenge.py b/backend/app/engine/nodes/challenge.py index ddec305e..00f0c22d 100644 --- a/backend/app/engine/nodes/challenge.py +++ b/backend/app/engine/nodes/challenge.py @@ -393,11 +393,13 @@ def _masked_rationale(outcome: VerifyOutcome, secrets: Any, *, given: Iterable[s 换的办法同审批质询的回答(_scrubbed_source):原文里的值先换掉,再照原来的目录重新组装理由文档;引用到的格子按 遮好的快照渲染(文字格里夹着的遮罩值换掉那一截)。理由里认不出遮罩值的原样交回。""" from app.engine.evidence import compose_doc - from app.engine.masking import _masked_view, _scrubbed_source + from app.engine.masking import _masked_view, _scrubbed_source, _Secrets doc, result = outcome.rationale, outcome.result - pool = secrets.copy() + # agent 查询时遮掉的值先收(masked_values 里已经排在报告的值前面):报告的值收满了上限 SECRETS_MAX,也不丢 agent 查到的 + pool = _Secrets() pool.merge(outcome.masked_values) + pool.merge(secrets) pool.protect(given) if result is None or not isinstance(doc, dict) or not pool.strong: return outcome diff --git a/backend/app/engine/verify_agent.py b/backend/app/engine/verify_agent.py index 9902ab0c..fd7b3bf8 100644 --- a/backend/app/engine/verify_agent.py +++ b/backend/app/engine/verify_agent.py @@ -7,12 +7,14 @@ (tools/datasource.py)。这一路没有审批关卡,不能指望人在半路拦下一条 DELETE。 **先遮再交。** 建查询工具时还打开遮罩开关(tools/datasource 的 mask_results):数据源设了遮罩时,取自遮罩列的格子 -(按 SQL 认,起别名、套函数、经子查询或 WITH 转一道、SELECT *、表值函数拆开的都算)在交给模型、落进查询快照之前就 -换成「已遮罩」,别的格子按值换:这份结果里的遮罩值、前面的查询遮掉的、调用方交进来的(masked_values:审批质询、 -反方审稿收的报告快照里的值),文字格里夹着的换那一截(engine/masking.py 的 mask_query_result);查询失败时报错原话 -里的值也换。模型看不到遮罩值,就不会把它写进回答、理由、SQL;以前是事后按值认,每一轮复核都找到新的漏法。快照里的 -masked_at_query 记下查询时就遮了的列,格子引用指到这几列、指到按值换掉的格子的都不采信。遮掉的值只在内存里交回 -调用方(VerifyOutcome.masked_values,审批质询、反方审稿按值遮的第二道防线用),不进快照、不交给模型、不进 as_dict。 +(按 SQL 认,起别名、套函数、经子查询或 WITH 转一道、SELECT *(展开的列换了名字的也算)、表值函数拆开的都算)在交给 +模型、落进查询快照之前就换成「已遮罩」,别的格子按值换:这份结果里的遮罩值、前面的查询遮掉的、调用方交进来的 +(masked_values:审批质询、反方审稿收的报告快照里的值),文字格里夹着的换那一截(engine/masking.py 的 +mask_query_result);查询失败时报错原话里的值也换。模型看不到遮罩值,就不会把它写进回答、理由、SQL;以前是事后 +按值认,每一轮复核都找到新的漏法。快照里的 masked_at_query 记下查询时就遮了的列,格子引用指到这几列、指到按值换掉 +的格子的都不采信。遮掉的值只在内存里交回调用方(VerifyOutcome.masked_values,审批质询、反方审稿按值遮的第二道防线 +用),不进快照、不交给模型、不进 as_dict。调用方交进来的和查询时遮掉的分开记(masking.SeenValues),合起来时 agent +查到的在前:合起来到了上限 SECRETS_MAX,挤掉的是报告里的值。 **数从快照取。** 每次成功的查询是一份查询快照(工具自己落的 put_json),按这个 agent 自己的顺序编 Q1、Q2…。 最后交结果时,调用方结构里的数值字段一律改成格子引用 {alias, row, column},由代码按引用从快照里取值:模型自己 @@ -67,7 +69,7 @@ from app.engine import spend from app.engine.evidence import build_catalog, cell_eid, compose_doc, query_entry_fields, schema_ledger_entry from app.engine.expressions import CellError, column_kind, locate_cell, numeric_text from app.engine.judge import FALLBACK_MAX_TOKENS, MAX_TOKENS, _json_in, _tokens, _usd -from app.engine.masking import _masked_word, _Secrets +from app.engine.masking import SeenValues, _masked_word, _Secrets from app.engine.nodes.llm import _DIGEST_BUDGET, _digest, _split_structured, _typed, explain_model_error from app.engine.state import message_text from app.engine.toolcalls import ToolTimeout, limit_fields, limit_of, run_bounded @@ -160,8 +162,9 @@ class VerifyOutcome: output_tokens, total_tokens, cost_usd, calls, unpriced_calls}。limit:触顶原因(LIMITS 之一),没触顶为 None。 answered:模型是不是自己停下来的(没被上限掐断,也没因为输出到了上限被截断)。被截断的一律没有结果,error 是 TRUNCATED_ERROR。 - masked_values:查询时在工具边界遮掉的值(engine/masking 的 _Secrets),连同调用方交进来的。只在内存里交给调用方按值遮 - (审批质询的回答、反方审稿的理由),不进 as_dict、不往外交。 + masked_values:查询时在工具边界遮掉的值(engine/masking 的 _Secrets),连同调用方交进来的:agent 查到的在前,合起来 + 到了上限 SECRETS_MAX 挤掉的是调用方的(SeenValues.pooled)。只在内存里交给调用方按值遮(审批质询的回答、反方审稿的 + 理由),不进 as_dict、不往外交。 """ result: dict[str, Any] | None @@ -670,9 +673,9 @@ class _State: queries: list[dict[str, Any]] = field(default_factory=list) ledger: list[dict[str, Any]] = field(default_factory=list) model: str | None = None - #: 查询工具在工具边界遮掉的值(经 ToolContext.extra 并进来,只在内存里),起头是调用方交进来的。后面的查询按它们 - #: 按值遮 - seen: _Secrets = field(default_factory=_Secrets) + #: 认出的遮罩值(只在内存里):调用方交进来的,加上查询工具在工具边界遮掉的(经 ToolContext.extra 并进来)。后面的 + #: 查询按它们按值遮,最后合起来交回调用方(run_verify_agent),agent 查到的在前 + seen: SeenValues = field(default_factory=SeenValues) def outcome(self, **fields: Any) -> VerifyOutcome: usage = self.usage @@ -682,8 +685,7 @@ class _State: for q in self.queries] base: dict[str, Any] = {"result": None, "cells": [], "queries": shown, "rationale": None, "usage": usage, "notes": self.notes, "model": self.model, - "duration_ms": int((time.monotonic() - self.started) * 1000), - "masked_values": self.seen} + "duration_ms": int((time.monotonic() - self.started) * 1000)} return VerifyOutcome(**{**base, **fields}) @@ -720,20 +722,25 @@ async def run_verify_agent( data_versions;不在运行里传 Versions.CURRENT;缺省 UNSET 时上传源一律报错,不拿当前版本顶上)。 - masked_values:调用方已经认出的遮罩值(engine/masking 的 _Secrets,审批质询、反方审稿收的报告快照里的值)。 查询工具从第一条查询起就按它们按值遮:agent 查到的备注、别的表里写着同一个号码时,模型看不到。随查询时遮掉的值 - 一起交回(VerifyOutcome.masked_values)。 + 一起交回(VerifyOutcome.masked_values)。不改它。 """ - state = _State(started=time.monotonic()) + state = _State(started=time.monotonic(), seen=SeenValues( + given=masked_values if masked_values is not None else _Secrets())) try: - if masked_values is not None: - # 上万个值并一遍要几十毫秒,放进线程,别挡住同一进程里别的运行 - await asyncio.to_thread(state.seen.merge, masked_values) - return await _verify(state, spec=spec, system=system, task=task, sources=sources, - output_schema=output_schema, max_queries=max_queries, timeout_s=timeout_s, - max_cost_usd=max_cost_usd, daily_max_usd=daily_max_usd, max_steps=max_steps, emit=emit, - run_id=run_id, node_id=node_id, data_versions=data_versions) + out = await _verify(state, spec=spec, system=system, task=task, sources=sources, + output_schema=output_schema, max_queries=max_queries, timeout_s=timeout_s, + max_cost_usd=max_cost_usd, daily_max_usd=daily_max_usd, max_steps=max_steps, emit=emit, + run_id=run_id, node_id=node_id, data_versions=data_versions) except Exception as e: # noqa: BLE001 - 永远不抛:库写不进去、建工具出错这类意外也落在返回值里 logger.warning("核查出错", exc_info=True) - return state.outcome(error=f"核查出错({describe_exception(e)})") + out = state.outcome(error=f"核查出错({describe_exception(e)})") + try: + # agent 查到的在前、调用方交进来的在后:上万个值合一遍要几十毫秒,放进线程,别挡住同一进程里别的运行 + pooled = await asyncio.to_thread(state.seen.pooled) + except Exception: # noqa: BLE001 - 永远不抛:合不起来就只交调用方交进来的(agent 的查询结果在工具边界已经遮过) + logger.warning("核查遮掉的值合并出错", exc_info=True) + pooled = state.seen.given + return replace(out, masked_values=pooled) async def _verify(state: _State, *, spec: ModelSpec, system: str, task: str, sources: Iterable[str], diff --git a/backend/app/tools/datasource.py b/backend/app/tools/datasource.py index 7f2a7e39..1cf26969 100644 --- a/backend/app/tools/datasource.py +++ b/backend/app/tools/datasource.py @@ -53,8 +53,9 @@ RUN_VERSIONS_KEY = "data_versions" DATA_UNAVAILABLE = "查询失败(数据不可用):" #: 运行中途补固定一个源时发的那条日志事件的 code(_pin_late) LATE_PIN_CODE = "data_version_pinned" -#: ToolContext.extra 里放「查询时遮掉的值」的键(engine/masking 的 _Secrets):核查 agent(engine/verify_agent.py)放进去, -#: 打开了遮罩开关的查询工具把每次遮掉的值并进去。只在内存里给调用方按值遮用,不进快照、不交给模型 +#: ToolContext.extra 里放「认出的遮罩值」的键(engine/masking 的 SeenValues:调用方交进来的,加上查询时遮掉的):核查 +#: agent(engine/verify_agent.py)放进去,打开了遮罩开关的查询工具按它们按值遮,再把每次遮掉的值并进去。只在内存里给 +#: 调用方按值遮用,不进快照、不交给模型 MASKED_VALUES_KEY = "masked_values" #: 每次查询之后基于数据目录的 SQL 检查最多等这么久。正常一条 SQL 几毫秒;等不到就放弃这一次检查, @@ -498,9 +499,11 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False, return f"SQL 被拒绝:{FORCED_READONLY}" limits = QueryLimits(timeout_seconds=seconds, **({"max_rows": int(limit)} if limit else {})) mask = masked_columns(source.options) - # 此前认出的遮罩值:交给线程的是一份拷贝,线程里读的时候这边不会再往里并 + # 此前认出的遮罩值合起来的另一份(agent 前面的查询遮掉的在前,调用方交进来的在后,见 masking.SeenValues):线程里 + # 读的时候这边不会再往里并。上万个值合一遍要几十毫秒,放进线程 sink = ctx.extra.get(MASKED_VALUES_KEY) if mask_results else None - known = sink.copy() if sink is not None and sink.strong else None + known = await asyncio.to_thread(sink.pooled) if sink is not None else None + known = known if known is not None and known.strong else None try: result = await run_query(source, sql, limits=limits) except SqlRejected as e: @@ -538,8 +541,8 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False, if mask or done.payload is not payload: payload, masked_now = done.payload, done.columns if sink is not None: - # 上万个值并一遍要几十毫秒(免遮的文字要重新过一遍),也放进线程 - await asyncio.to_thread(sink.merge, done.secrets) + # 并进 agent 查到的那一份(不和调用方交进来的混在一起,上限收满了也不丢),也放进线程 + await asyncio.to_thread(sink.add, done.secrets) # 对照数据目录检查这条 SQL(一对多关联后重复计算、存量跨期求和……)。结果交给模型(Agent 据此改写), # 也随查询快照存下:证据面板看得到,口径卡据此给指标标「存疑」、出具降档。没查出问题不加这个键, # 快照和以前一字不差 diff --git a/backend/tests/test_challenge_node.py b/backend/tests/test_challenge_node.py index 441b205a..22ded299 100644 --- a/backend/tests/test_challenge_node.py +++ b/backend/tests/test_challenge_node.py @@ -1405,3 +1405,26 @@ async def test_the_node_refuses_the_same_bad_settings_at_run_time(): with pytest.raises(NodeError) as caught: await challenge_mod.run_challenge({"nodes": {}}, ctx) assert str(caught.value) == message + + +def test_values_the_agent_found_are_masked_in_the_rationale_when_the_report_fills_the_cap(monkeypatch): + """复核第 2 轮:报告快照里的遮罩值收满了上限 SECRETS_MAX,反方 agent 查询时遮掉的号码照样认得出(agent 查到的先收)。 + 以前先收报告的值、满了就停,agent 查到的号码写进理由原样存下。""" + from app.engine import masking + from app.engine.evidence import compose_doc + from app.engine.verify_agent import VerifyOutcome + + monkeypatch.setattr(masking, "SECRETS_MAX", 40) + secrets = masking._Secrets() + secrets.add({"columns": ["p"], "rows": [[str(13800000000 + i)] for i in range(40)]}, {"p"}) + assert len(secrets.weak) == masking.SECRETS_MAX + found = masking._Secrets() + found.add({"columns": ["p"], "rows": [["19914742785"]]}, {"p"}) + # 和 run_verify_agent 交回的一样:agent 查到的在前,调用方交进去的在后 + seen = masking.SeenValues(given=secrets) + seen.add(found) + doc = compose_doc("号码 199-1474-2785 的客户在华南。", {}, node_id="chal", run_id="r1") + outcome = VerifyOutcome(result={"verdict": "stands", "rationale": doc["markdown"]}, cells=[], queries=[], + rationale=doc, usage={}, masked_values=seen.pooled()) + masked = challenge_mod._masked_rationale(outcome, secrets, given=[]) + assert masked.result["rationale"] == "号码 已遮罩 的客户在华南。" diff --git a/backend/tests/test_inquiry_api.py b/backend/tests/test_inquiry_api.py index 275d612a..2f09bbd5 100644 --- a/backend/tests/test_inquiry_api.py +++ b/backend/tests/test_inquiry_api.py @@ -933,6 +933,49 @@ async def test_the_inquiry_agent_sees_and_stores_its_queries_masked(monkeypatch, assert content == artifact_store.load(query["artifact"]) +@pytest.mark.parametrize("sql, rows", [ + pytest.param("SELECT '合计' AS 区域, SUM(amount) AS 金额 FROM sales UNION ALL SELECT * FROM " + "(SELECT buyer_phone, amount FROM sales ORDER BY amount DESC LIMIT 2)", + [["已遮罩", 2800], ["已遮罩", 1200], ["已遮罩", 900]], id="UNION 后一支的 *"), + pytest.param("WITH t(a, b, c, d, e) AS (SELECT * FROM sales) SELECT c AS 区域, d AS 金额 FROM t " + "ORDER BY d DESC LIMIT 2", [["已遮罩", 1200], ["已遮罩", 900]], id="CTE 列名表套 *"), +]) +async def test_a_star_renaming_the_masked_column_never_shows(monkeypatch, tmp_path, client, sql, rows): + """复核第 2 轮的 blocker:* 把手机号带进了别的列名(UNION 后一支的 *,结果的列名取第一支的「区域」;CTE 的列名表把 + SELECT * 展开的列换了名字)。以前交给模型的工具结果里就有原值,回答「金额最大的一笔来自 13900000001」照常交出,事件、 + 质询工件、通用工件接口都给原值。现在 * 展开落到的那几列查询时就遮掉:只展开一个子查询的按位置只遮号码那一列;展开 + 数据源里的表的按查询当时的表结构数出号码排在第几,也只遮那一列。""" + name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=SALES) + plays = {"最大": Play([sql], "金额最大的一笔来自 [[v:Q1.r1.区域]]。[[see:Q1]]")} + Models(monkeypatch, name, writer="华东销售额 [[v:Q1.r0.销售额]]。[[see:Q1]]", plays=plays) + seen: list[str] = [] + answering = MockChatModel._decide + + def spy(model, messages): + seen.extend(str(m.content) for m in messages if isinstance(m, ToolMessage)) + return answering(model, messages) + + monkeypatch.setattr(MockChatModel, "_decide", spy) + try: + paused = await pause(name, graph=fetching(name, QUERY)) + [approval] = await pending(paused.id) + r = await ask(client, paused.id, report="write", unit=None, question="金额最大的一笔来自哪里?", + approval_id=approval.id) + assert r.status_code == 200, r.text + item = r.json()["inquiry"] + assert (item["answer"]["ok"], item["answer"]["text"]) == (True, "金额最大的一笔来自 已遮罩。") + assert seen and all("1390000000" not in s for s in seen), "交给模型的查询结果里就是遮好的" + [query] = item["queries"] + for url in (f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}", + f"/api/artifacts/{query['artifact']}"): + got = await client.get(url) + assert got.status_code == 200 and got.json()["content"]["rows"] == rows and "1390000000" not in got.text + [event] = await events(paused.id, "approval.inquiry") + assert "1390000000" not in json.dumps(event.data, ensure_ascii=False) + finally: + await engines.invalidate(source_id) + + async def test_a_masked_column_under_another_name_is_masked_everywhere(monkeypatch, masked_shop, client): """遮罩列起了中文别名(本仓库查询的惯例):回答里引用到的格子、事件、回答文档、质询工件里都不露出原值。 报告自己的查询没选手机号,原值哪里都认不出来,只能按 SQL 认出「手机号」就是 buyer_phone。""" @@ -1552,6 +1595,40 @@ async def test_a_number_from_the_agents_other_query_is_masked_when_the_number_ca await engines.invalidate(source_id) +async def test_a_number_the_agent_found_is_masked_even_when_report_values_fill_the_cap(monkeypatch, tmp_path, client): + """复核第 2 轮:报告快照里的遮罩值收满了上限(这里把 SECRETS_MAX 调成报告正好收满的大小)。agent 先查出华东的号码 + (查询时遮掉),再查回访备注:备注里的号码照样在交给模型之前换掉,回答、事件、质询工件里都没有。以前开始回答时先并 + 报告的值、收满就停,agent 查到的号码并不进去:备注原样交给模型、原样落库,回答引用这一格时第二道防线也认不出,原值 + 写进了运行记录。""" + from app.engine import masking + + name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CALLBACKS) + report_sql = "SELECT buyer_phone, amount AS 金额 FROM sales WHERE region <> '华东' ORDER BY id" + plays = {"回访": Play([PHONES_EAST, NOTES], "华东的回访备注是 [[v:Q2.r0.备注]]。[[see:Q2]]")} + Models(monkeypatch, name, writer="华南一笔金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays) + try: + paused = await pause(name, graph=fetching(name, report_sql)) + [approval] = await pending(paused.id) + report = [artifact_store.load(e.data["query_artifact"]) for e in await events(paused.id, "tool.end") + if e.data.get("query_artifact")] + full = masking._Secrets() + full.add(report[0], {"buyer_phone"}) + monkeypatch.setattr(masking, "SECRETS_MAX", len(full.weak)) + r = await ask(client, paused.id, report="write", unit=None, question="华东的回访说了什么?", + approval_id=approval.id) + assert r.status_code == 200, r.text + item = r.json()["inquiry"] + assert (item["answer"]["ok"], item["answer"]["text"]) == (True, "华东的回访备注是 客户已遮罩来电改址。") + phones, notes = item["queries"] + assert artifact_store.load(notes["artifact"])["rows"] == [["华东", "客户已遮罩来电改址"], ["华南", "例行回访"]] + [event] = await events(paused.id, "approval.inquiry") + assert "1390000000" not in json.dumps(event.data, ensure_ascii=False) + got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{notes['artifact']}") + assert got.status_code == 200 and "1390000000" not in got.text + finally: + await engines.invalidate(source_id) + + async def test_a_number_from_the_agents_other_query_is_masked_when_the_number_came_last(monkeypatch, tmp_path, client): """先查备注、后查号码:查备注时还不知道里面写着遮罩值,原样交给模型、原样存下;回答引用这一格,渲染出来才露,整份 不交。回答时按内存里的值换掉的位置记进事件(不含值),事后取质询工件照着再换一遍:质询工件里那一截是「已遮罩」, diff --git a/backend/tests/test_masking_star_renames.py b/backend/tests/test_masking_star_renames.py new file mode 100644 index 00000000..36c93230 --- /dev/null +++ b/backend/tests/test_masking_star_renames.py @@ -0,0 +1,251 @@ +"""按 SQL 认取自遮罩列的结果列(engine/masking.py 的 _derived_hidden):* 把遮罩列带进了别的列名(复核第 2 轮的 blocker)。 + +* 展开的列按位置落进结果:WITH t(a, b, …) AS (SELECT * FROM sales) 的列名表、FROM 里子查询的列名表会给它们换上别的 +名字,UNION 后一支的 * 落进第一支起的列名里,标量子查询里的 * 则把值交给外层的表达式。以前 * 本身不算取自哪一列,这些 +写法的原值原样交给核查 agent、落进快照,审批质询按值遮的第二道防线也认不出。现在 * 带出的值取自遮罩列时,它展开落到的 +那几列都算:只展开一个来源、那里的每一列都数得清的(子查询按选出项,数据源里的表按查询当时的表结构),按位置只算取自 +遮罩列的那几列;数不清的(没有表结构、列数对不上)整段都算(多遮)。不改名的 *(单独一条 SELECT 的 *、UNION 第一支的 *) +照旧按名字认;* 只带出没碰遮罩列的值、EXISTS 里的 * 都不算。 + +查询都在内存里的 SQLite 上真跑一遍,看遮完的结果里还有没有原值。数据都是虚构的:区域、手机号、姓名、金额。 +""" +from __future__ import annotations + +import json +import sqlite3 + +import pytest + +from app.db import models # noqa: F401 - 建表(conftest)之前先登记模型:表结构快照要存进工件库 +from app.engine import masking + +MASKED_WORD = "已遮罩" +#: 设了遮罩的列 +MASKED = ["buyer_phone", "buyer_name"] +#: 遮罩列里的值:遮完的结果里不能出现 +SECRETS = ("1390000000", "张伟", "李娜", "王芳") +SCRIPT = ("CREATE TABLE sales (id INTEGER, region TEXT, buyer_phone TEXT, buyer_name TEXT, amount INTEGER, note TEXT);" + "INSERT INTO sales VALUES (1, '华东', '13900000001', '张伟', 1200, '改址')," + " (2, '华南', '13900000002', '李娜', 300, '回访'), (3, '华北', '13900000003', '王芳', 900, '');" + # 只有一列号码的表:标量子查询、UNION 后一支的 * 直接展开它 + "CREATE TABLE phones (buyer_phone TEXT);" + "INSERT INTO phones VALUES ('13900000009');") + + +@pytest.fixture(scope="module") +def db(): + conn = sqlite3.connect(":memory:") + conn.executescript(SCRIPT) + yield conn + conn.close() + + +def masked_run(db, sql: str) -> masking.QueryMask: + """真跑一遍 SQL,按工具边界的办法遮(mask_query_result)。""" + cursor = db.execute(sql) + columns = [d[0] for d in cursor.description] + rows = [list(r) for r in cursor.fetchall()] + payload = {"columns": columns, "rows": rows, "row_count": len(rows), "truncated": False, "sql": sql} + return masking.mask_query_result(payload, MASKED, "sqlite") + + +def _leaks(done: masking.QueryMask) -> bool: + text = json.dumps(done.payload["rows"], ensure_ascii=False) + return any(s in text for s in SECRETS) + + +@pytest.mark.parametrize("sql, hidden, shown", [ + pytest.param("WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT * FROM t", {"c", "d"}, set(), + id="CTE 列名表套 * 展开表"), + pytest.param("WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT c AS 号码, e AS 金额 FROM t " + "ORDER BY e DESC LIMIT 1", {"号码"}, set(), id="CTE 列名表套 * 再改名"), + pytest.param("SELECT 0 AS a, '' AS b, '' AS c, '' AS d, 0 AS e, '' AS f UNION ALL SELECT * FROM sales", + {"c", "d"}, set(), id="UNION 后一支的 * 展开表"), + pytest.param("SELECT region FROM sales UNION ALL SELECT * FROM phones", {"region"}, set(), + id="UNION 后一支的 * 展开只有号码的表"), + pytest.param("WITH a AS (SELECT buyer_phone FROM sales), b(x) AS (SELECT * FROM a) SELECT x FROM b", {"x"}, + set(), id="列名表套 * 展开另一段 CTE"), + pytest.param("SELECT r FROM (SELECT region AS r FROM sales UNION ALL " + "SELECT * FROM (SELECT buyer_phone FROM sales))", {"r"}, set(), id="子查询里 UNION 后一支的 *"), + pytest.param("SELECT * FROM (SELECT region FROM sales UNION ALL SELECT * FROM (SELECT buyer_phone FROM sales))", + {"region"}, set(), id="外层 * 取子查询里 UNION 后一支的 *"), + pytest.param("WITH t(k) AS (SELECT * FROM (SELECT buyer_name FROM sales)) SELECT k AS 名 FROM t", {"名"}, set(), + id="列名表套 * 展开子查询"), + pytest.param("SELECT region FROM sales UNION ALL SELECT s.* FROM (SELECT buyer_name FROM sales) s", {"region"}, + set(), id="UNION 后一支的 t.*"), + pytest.param("SELECT region FROM sales UNION SELECT * FROM (SELECT buyer_phone FROM sales)", {"region"}, set(), + id="UNION 去重"), + # 只展开一个派生来源、每一列都数得清:按位置只遮取自遮罩列的那一列,金额照常 + pytest.param("SELECT '合计' AS 区域, SUM(amount) AS 金额 FROM sales UNION ALL SELECT * FROM " + "(SELECT buyer_phone, amount FROM sales ORDER BY amount DESC LIMIT 2)", {"区域"}, {"金额"}, + id="合计行 UNION 后一支的 * 按位置"), + pytest.param("SELECT region AS 区域, amount AS 金额 FROM sales UNION ALL SELECT * FROM " + "(SELECT * FROM (SELECT buyer_name, amount FROM sales))", {"区域"}, {"金额"}, + id="套了两层的 * 按位置"), + pytest.param("WITH t(r, x) AS (SELECT * FROM (SELECT buyer_phone, amount FROM sales)) SELECT r AS 号码, x AS 金额 " + "FROM t", {"号码"}, {"金额"}, id="列名表套 * 按位置"), + pytest.param("SELECT (SELECT * FROM (SELECT buyer_phone FROM sales) LIMIT 1) AS x", {"x"}, set(), + id="标量子查询里的 * 展开子查询"), + pytest.param("SELECT region, (SELECT * FROM (SELECT s2.buyer_name FROM sales s2 WHERE s2.id = s.id)) AS 名 " + "FROM sales s", {"名"}, {"region"}, id="相关子查询里的 *"), + pytest.param("SELECT (SELECT * FROM phones LIMIT 1) AS x", {"x"}, set(), id="标量子查询里的 * 展开表"), +]) +def test_a_star_that_renames_a_masked_column_is_masked(db, sql, hidden, shown): + """* 把遮罩列的值带进了别的列名:这几列整列遮,原值不交出去。按位置认得出的,别的列照常。""" + done = masked_run(db, sql) + columns = {c.lower() for c in done.columns} + assert hidden <= columns, done.columns + assert not shown & columns, done.columns + assert not _leaks(done), done.payload["rows"] + assert done.secrets.strong, "遮掉的值交回调用方(后面的查询、回答里按值遮)" + + +@pytest.mark.parametrize("sql", [ + "SELECT * FROM (SELECT region, amount FROM sales) t", + "WITH t AS (SELECT * FROM sales) SELECT region, SUM(amount) AS 合计 FROM t GROUP BY region", + "SELECT region AS 区域 FROM sales UNION ALL SELECT * FROM (SELECT region FROM sales)", + "WITH t(a, b) AS (SELECT * FROM (SELECT region, amount FROM sales)) SELECT a, SUM(b) AS 合计 FROM t GROUP BY a", + "SELECT region, EXISTS(SELECT * FROM sales s2 WHERE s2.id = s.id AND s2.amount > 500) AS 大单 FROM sales s", + "SELECT '合计' AS 区域, SUM(amount) AS 金额 FROM sales UNION ALL SELECT * FROM " + "(SELECT region, SUM(amount) FROM sales GROUP BY region)", +]) +def test_a_star_that_brings_no_masked_values_stays_visible(db, sql): + """* 只带出没碰遮罩列的值、EXISTS 里的 *(只说有没有):一列都不遮,结果原样。""" + done = masked_run(db, sql) + assert done.columns == [] and MASKED_WORD not in json.dumps(done.payload, ensure_ascii=False) + + +def test_a_plain_star_still_masks_by_name(db): + """不改名的 *(单独一条 SELECT)照旧只遮名字对上遮罩列的那几列。""" + done = masked_run(db, "SELECT * FROM sales") + assert done.columns == ["buyer_phone", "buyer_name"] and not _leaks(done) + + +@pytest.mark.parametrize("sql, columns, hidden", [ + pytest.param("SELECT x FROM (SELECT * FROM (SELECT buyer_phone FROM sales) s) AS q(x)", ["x"], {"x"}, + id="子查询列名表套 *"), + pytest.param("SELECT l.x FROM sales s CROSS JOIN LATERAL (SELECT * FROM (SELECT s.buyer_phone) z) AS l(x)", ["x"], + {"x"}, id="LATERAL 列名表套 *"), + pytest.param("SELECT ARRAY(SELECT * FROM (SELECT buyer_phone FROM sales) s) AS arr", ["arr"], {"arr"}, + id="ARRAY 子查询里的 *"), +]) +def test_a_star_renaming_in_postgres_is_masked(sql, columns, hidden): + snapshot = {"source": "shop", "sql": sql, "columns": columns, "rows": []} + assert hidden <= masking._derived_hidden(snapshot, set(MASKED), "postgres")[0] + + +def test_a_metric_summed_over_a_star_renamed_masked_column_is_not_sent_to_blind_check(): + """双盲复算按同一套认法挡下取自遮罩列的指标(derived_columns):遮罩的金额经 CTE 列名表套 * 改名之后汇总,照样认得出, + 区域照常。""" + sql = "WITH q(r, x) AS (SELECT * FROM (SELECT region, amount FROM sales)) SELECT r AS 区域, SUM(x) AS gmv FROM q " \ + "GROUP BY r" + snapshot = {"source": "shop", "sql": sql, "columns": ["区域", "gmv"], "rows": []} + assert masking.derived_columns(snapshot, {"amount"}, "sqlite") == {"gmv"} + + +def test_a_recursive_cte_with_a_star_is_read_without_looping(db): + """递归的 CTE 里 * 展开它自己:照样认得出、不会绕不出来。""" + sql = ("WITH RECURSIVE r(n, p) AS (SELECT 1, buyer_phone FROM sales WHERE id = 1 UNION ALL " + "SELECT n + 1, p FROM r WHERE n < 2) SELECT * FROM (SELECT * FROM r) t") + done = masked_run(db, sql) + assert "p" in {c.lower() for c in done.columns} and not _leaks(done) + loop = "WITH RECURSIVE r AS (SELECT buyer_phone AS p FROM sales UNION ALL SELECT * FROM r) SELECT p AS 号 FROM r" + snapshot = {"source": "shop", "sql": loop, "columns": ["号"], "rows": []} + assert "号" in masking._derived_hidden(snapshot, set(MASKED), "sqlite")[0] + + +def test_a_star_with_a_replace_is_not_counted_by_position(): + """* 带着 REPLACE / EXCLUDE 这类改动(DuckDB 一类的写法):展开成的列不再是来源那几列的原样,不按位置数,整段都算。 + 这里 REPLACE 把手机号换进了「amount」那一列。""" + sql = ("SELECT region, amount FROM sales UNION ALL SELECT * REPLACE (buyer_phone AS amount) FROM " + "(SELECT buyer_phone, amount FROM sales) t") + snapshot = {"source": "shop", "sql": sql, "columns": ["region", "amount"], "rows": []} + assert {"region", "amount"} <= masking._derived_hidden(snapshot, set(MASKED), "duckdb")[0] + + +# -------------------------------------------------------------------------- +# * 展开数据源里的表:按查询当时的表结构数位置 +# -------------------------------------------------------------------------- + + +#: 两个月的订单表列一样、次序一样(上传表格里的月度工作表常这样);三月那张的号码和区域换了次序 +MONTHS = ("CREATE TABLE jan (id INTEGER, region TEXT, buyer_phone TEXT, amount INTEGER);" + "CREATE TABLE feb (id INTEGER, region TEXT, buyer_phone TEXT, amount INTEGER);" + "CREATE TABLE mar (id INTEGER, buyer_phone TEXT, region TEXT, amount INTEGER);" + "INSERT INTO jan VALUES (1, '华东', '13900000001', 1200), (2, '华南', '13900000002', 300);" + "INSERT INTO feb VALUES (3, '华东', '13900000003', 900);" + "INSERT INTO mar VALUES (4, '13900000004', '华北', 400);") + + +@pytest.fixture(scope="module") +def months(): + conn = sqlite3.connect(":memory:") + conn.executescript(MONTHS) + yield conn + conn.close() + + +def schema_of(conn) -> dict: + """数据源的表结构(schema_cache 的样子):每张表的列按表里的次序。""" + names = [name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name")] + return {"tables": {name: {"columns": [{"name": row[1]} for row in conn.execute(f"PRAGMA table_info({name})")]} + for name in names}} + + +def run_months(conn, sql: str, schema: dict | None) -> masking.QueryMask: + cursor = conn.execute(sql) + payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql} + return masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=schema) + + +def test_a_union_of_tables_with_the_same_layout_is_masked_by_name_with_the_schema(months): + """UNION 的几支都是 SELECT * FROM 表、表的列一样:按表结构数得出号码落在哪一列,只遮那一列,汇总照常。没有表结构时 + 说不清号码排在第几,* 展开落到的那一段整段遮(多遮)。""" + schema = schema_of(months) + total = "SELECT region, SUM(amount) AS 合计 FROM (SELECT * FROM jan UNION ALL SELECT * FROM feb) t GROUP BY region" + done = run_months(months, total, schema) + assert done.columns == [] and done.payload["rows"] == [["华东", 2100], ["华南", 300]] + assert {"region", "合计"} <= set(run_months(months, total, None).columns) + both = run_months(months, "SELECT * FROM jan UNION ALL SELECT * FROM feb", schema) + assert both.columns == ["buyer_phone"] and not _leaks(both) + + +def test_a_union_of_tables_in_another_order_masks_where_the_values_land(months): + """三月那张表号码排在第二列:UNION 之后它落进第一支的「region」,按表结构认得出,这一列遮;金额照常。""" + schema = schema_of(months) + both = run_months(months, "SELECT * FROM jan UNION ALL SELECT * FROM mar", schema) + assert {"region", "buyer_phone"} <= set(both.columns) and "amount" not in both.columns and not _leaks(both) + total = "SELECT region AS 区域, SUM(amount) AS 合计 FROM (SELECT * FROM jan UNION ALL SELECT * FROM mar) t " \ + "GROUP BY region" + done = run_months(months, total, schema) + assert done.columns == ["区域"] and not _leaks(done) + + +def test_a_column_list_over_a_table_is_masked_by_position_with_the_schema(months): + """列名表套 * 展开数据源里的表:按表结构只遮号码换了名字的那一列。""" + done = run_months(months, "WITH t(a, b, c, d) AS (SELECT * FROM jan) SELECT * FROM t", schema_of(months)) + assert done.columns == ["c"] and [row[3] for row in done.payload["rows"]] == [1200, 300] + + +def test_a_schema_that_does_not_match_the_table_masks_the_whole_stretch(months): + """表结构里没有这张表,或者列数和实际展开的对不上(表结构过期了):照旧整段遮。""" + stale = schema_of(months) + stale["tables"]["feb"]["columns"] = stale["tables"]["feb"]["columns"][:3] + for schema in (stale, {"tables": {"jan": schema_of(months)["tables"]["jan"]}}): + both = run_months(months, "SELECT * FROM jan UNION ALL SELECT * FROM feb", schema) + assert set(both.columns) == {"id", "region", "buyer_phone", "amount"} and not _leaks(both) + + +async def test_a_snapshot_is_read_with_its_frozen_schema(months): + """事后按快照认(双盲复算挡指标、审批质询收值)用快照记着的表结构快照(schema_artifact),和查询当时工具边界的认法 + 一样:几张同样的月度表 UNION 起来汇总的金额不算取自遮罩列。""" + from app.core import artifact_store + + sql = "SELECT region, SUM(amount) AS 合计 FROM (SELECT * FROM jan UNION ALL SELECT * FROM feb) t GROUP BY region" + schema = await artifact_store.put_json({"source": "shop", **schema_of(months)}, kind="schema_snapshot", + run_id="r-star", node_id="n1") + snapshot = {"source": "shop", "sql": sql, "columns": ["region", "合计"], "rows": [], "schema_artifact": schema} + assert masking.derived_columns(snapshot, {"buyer_phone"}, "sqlite") == set() + assert masking._derived_hidden(snapshot, {"buyer_phone"}, "sqlite")[0] == {"buyer_phone"} + assert masking.derived_columns({**snapshot, "schema_artifact": "missing"}, {"buyer_phone"}, "sqlite") == { + "region", "合计"} diff --git a/backend/tests/test_masking_table_functions.py b/backend/tests/test_masking_table_functions.py index 272d3959..f37bb5e1 100644 --- a/backend/tests/test_masking_table_functions.py +++ b/backend/tests/test_masking_table_functions.py @@ -107,3 +107,55 @@ def test_a_masked_column_only_in_the_condition_of_an_aggregated_case_is_like_a_f ]) def test_a_case_returning_the_masked_value_or_not_aggregated_is_still_derived(sql): assert "total" in hidden_of(sql, ["total"]) + + +# -------------------------------------------------------------------------- +# 函数参数里的 JSON 箭头(复核第 2 轮) +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("kind, sql, columns, shown", [ + pytest.param("sqlite", "SELECT j.value AS v FROM sales, json_each(profile->'$.tels') AS j", ["v"], ["v"], + id="SQLite json_each 参数带箭头"), + pytest.param("postgres", "SELECT e AS 号 FROM sales, json_array_elements_text(profile->'tels') AS e", ["号"], + ["号"], id="Postgres json_array_elements_text"), + pytest.param("postgres", "SELECT e.value AS 号 FROM sales, jsonb_array_elements(profile->'tels') AS e", ["号"], + ["号"], id="Postgres jsonb_array_elements"), + pytest.param("postgres", "SELECT t.tel FROM sales, jsonb_to_recordset(profile->'contacts') AS t(tel text)", + ["tel"], ["tel"], id="Postgres jsonb_to_recordset 列名表"), + pytest.param("postgres", "SELECT k, v FROM sales, jsonb_each_text(profile->'contact') AS x(k, v)", ["k", "v"], + ["k", "v"], id="Postgres jsonb_each_text 列名表"), + pytest.param("postgres", "SELECT jsonb_array_elements_text(profile->'tels') AS 号 FROM sales", ["号"], ["号"], + id="Postgres 选出项里拆开"), + pytest.param("postgres", "SELECT jsonb_array_elements_text(profile -> 'contact' -> 'tels') AS 号 FROM sales", + ["号"], ["号"], id="Postgres 连着两个箭头"), + pytest.param("mysql", "SELECT JSON_UNQUOTE(profile->'$.tel') AS 电话 FROM sales", ["电话"], ["电话"], + id="MySQL JSON_UNQUOTE"), + pytest.param("mysql", "SELECT CONCAT(profile->'$.tel', '') AS 电话 FROM sales", ["电话"], ["电话"], + id="MySQL CONCAT"), +]) +def test_a_json_arrow_inside_function_arguments_is_read_as_the_masked_column(kind, sql, columns, shown): + """函数参数里没写表名的 col->'path' 被解析器当成以 col 为参数的 lambda(不是一列):以前认不出遮罩列,拆开、取出来的 + 值原样交给模型、落库。遮罩列是 JSON 时这几种是最自然的写法。写了表名、用 ->> 的本来就认得出。""" + assert {c.lower() for c in shown} <= hidden_of(sql, columns, kind) + + +def test_a_json_arrow_over_an_unmasked_column_stays_visible(): + """箭头取的不是遮罩列:一列都不遮。真的 lambda 的参数照列名收,只有和遮罩列同名时才多遮。""" + assert hidden_of("SELECT j.value AS 标签 FROM sales, json_each(tags->'$.list') AS j", ["标签"]) == MASKED + assert hidden_of("SELECT transform(scores, x -> x + 1) AS 分 FROM sales", ["分"], "postgres") == MASKED + + +def test_a_json_arrow_inside_json_each_is_masked_when_really_run(): + """SQLite 里真跑一遍:json_each(profile->'$.tels') 拆出来的号码在交给模型之前就遮掉。""" + import sqlite3 + + db = sqlite3.connect(":memory:") + db.executescript("CREATE TABLE sales (id INTEGER, profile TEXT);" + "INSERT INTO sales VALUES (1, '{\"tels\": [\"13900000001\", \"13900000002\"]}');") + sql = "SELECT j.value AS v FROM sales, json_each(profile->'$.tels') AS j" + cursor = db.execute(sql) + payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql} + db.close() + done = masking.mask_query_result(payload, ["profile"], "sqlite") + assert done.columns == ["v"] and done.payload["rows"] == [["已遮罩"], ["已遮罩"]] diff --git a/backend/tests/test_verify_agent_masking.py b/backend/tests/test_verify_agent_masking.py index 50725c47..ea6f2ab3 100644 --- a/backend/tests/test_verify_agent_masking.py +++ b/backend/tests/test_verify_agent_masking.py @@ -58,10 +58,10 @@ async def clean_spend(): yield -async def _source(tmp_path, options: dict | None) -> tuple[str, int]: +async def _source(tmp_path, options: dict | None, script: str = SALES) -> tuple[str, int]: path = tmp_path / f"{uuid.uuid4().hex[:8]}.db" db = sqlite3.connect(path) - db.executescript(SALES) + db.executescript(script) db.commit() db.close() name = f"shop_{uuid.uuid4().hex[:8]}" @@ -385,3 +385,108 @@ def test_labels_written_in_this_query_are_not_masked_by_values_from_earlier_ones payload = {"sql": "SELECT CASE WHEN amount > 1000 THEN '大单' ELSE '小单' END AS 档, COUNT(*) AS 笔数 FROM sales " "GROUP BY 1", "columns": ["档", "笔数"], "rows": [["大单", 1], ["小单", 3]]} assert masking.mask_query_result(payload, ["amount"], "sqlite", known).payload["rows"] == [["大单", 1], ["小单", 3]] + + +# -------------------------------------------------------------------------- +# * 把遮罩列带进别的列名、函数参数里的 JSON 箭头(复核第 2 轮) +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("sql, masked", [ + pytest.param("SELECT region, amount FROM sales WHERE id < 0 UNION ALL " + "SELECT * FROM (SELECT buyer_phone, amount FROM sales ORDER BY id)", ["region"], + id="UNION 后一支的 *"), + pytest.param("WITH t(k, v) AS (SELECT * FROM (SELECT buyer_name, amount FROM sales ORDER BY id)) " + "SELECT k AS 姓名, v AS amount FROM t", ["姓名"], id="CTE 列名表套 *"), + pytest.param("SELECT (SELECT * FROM (SELECT s2.buyer_phone FROM sales s2 WHERE s2.id = s.id)) AS 号码, s.amount " + "FROM sales s ORDER BY s.id", ["号码"], id="标量子查询里的 *"), +]) +async def test_a_star_renaming_a_masked_column_is_masked_before_the_model_sees_it(monkeypatch, masked_shop, sql, + masked): + """* 把遮罩列的值带进了别的列名(UNION 后一支、写了列名表的 CTE、标量子查询):交给模型的、落库的、通用接口给的都是 + 「已遮罩」,金额照常。以前 * 本身不算取自哪一列,原值原样交给模型、落进快照。""" + name = masked_shop + script = Script(monkeypatch, [q(name, sql, "c1")], + {"value": {"alias": "Q1", "row": 0, "column": "amount"}, "rationale": "见 [[see:Q1]]。"}) + out, _ = await ask(name) + assert out.error is None and out.result["value"] == 1200 + seen = script.tool_results()["c1"] + assert MASKED in seen and _clean(seen), seen + snapshot = artifact_store.load(out.queries[0]["artifact"]) + assert snapshot["masked_at_query"] == masked and set(_cells(snapshot, masked[0])) == {MASKED} + assert _cells(snapshot, "amount") == [1200, 300, 900, 400] + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + got = await client.get(f"/api/artifacts/{out.queries[0]['artifact']}") + assert got.status_code == 200 and got.json()["content"] == snapshot and _clean(got.text) + + +@pytest.mark.parametrize("sql, masked", [ + pytest.param("SELECT c.id AS 编号, j.value AS 值 FROM customers c, json_each(profile->'$') AS j " + "ORDER BY c.id, j.key", ["值"], id="json_each 参数带箭头"), + pytest.param("SELECT id AS 编号, upper(profile->'$.name') AS 姓名 FROM customers ORDER BY id", ["姓名"], + id="函数参数里的箭头"), +]) +async def test_a_json_arrow_over_the_masked_column_is_masked_before_the_model_sees_it(monkeypatch, json_shop, sql, + masked): + """遮罩的是 JSON 列:函数参数里没写表名的 profile->'…' 被解析器当成 lambda,以前认不出,拆开、取出来的号码和姓名原样 + 交给模型、落库。""" + name = json_shop + script = Script(monkeypatch, [q(name, sql, "c1")], + {"value": {"alias": "Q1", "row": 0, "column": "编号"}, "rationale": "x"}) + out, _ = await ask(name) + assert out.result["value"] == 1 and _clean(script.tool_results()["c1"]) + snapshot = artifact_store.load(out.queries[0]["artifact"]) + assert snapshot["masked_at_query"] == masked and set(_cells(snapshot, masked[0])) == {MASKED} + assert _clean(json.dumps(snapshot, ensure_ascii=False)) + + +# -------------------------------------------------------------------------- +# 遮罩值的上限:agent 自己查到的值先收(复核第 2 轮) +# -------------------------------------------------------------------------- + + +async def test_values_the_agent_finds_survive_caller_values_that_fill_the_cap(monkeypatch, masked_shop): + """调用方交进来的值(报告快照里的)收满了上限 SECRETS_MAX:agent 自己查到的号码照样收,后面查到的备注按它遮,交回 + 调用方的 masked_values 里也认得出。以前先并调用方的值、满了就停,agent 查到的一个也并不进去:后面的备注原样交给 + 模型、原样落库,回答时按值遮也认不出。""" + monkeypatch.setattr(masking, "SECRETS_MAX", 40) + name = masked_shop + given = masking._Secrets() + given.add({"columns": ["buyer_phone"], "rows": [[str(13800000000 + i)] for i in range(40)]}, {"buyer_phone"}) + assert len(given.weak) == masking.SECRETS_MAX and not given.hit("13900000001") + steps = [q(name, "SELECT buyer_phone, amount AS total FROM sales ORDER BY id", "c1"), + q(name, "SELECT id, note FROM sales ORDER BY id", "c2")] + script = Script(monkeypatch, steps, {"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "x"}) + out, _ = await ask(name, masked_values=given) + assert out.result["value"] == 1200 + assert _clean(script.tool_results()["c2"]), script.tool_results()["c2"] + notes = artifact_store.load(out.queries[1]["artifact"]) + assert _cells(notes, "note") == [f"客户{MASKED}来电改址", "例行回访", f"回访:联系 {MASKED}", ""] + assert out.masked_values.hit("13900000001") and out.masked_values.hit("13900000003") + assert len(out.masked_values.weak) <= masking.SECRETS_MAX, "合起来照样不超过上限:挤掉的是调用方的值" + + +#: 两个月的订单表,列一样、次序一样(上传表格里的月度工作表常这样) +MONTHS = ("CREATE TABLE jan (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, amount INTEGER);" + "CREATE TABLE feb (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, amount INTEGER);" + "INSERT INTO jan VALUES (1, '华东', '13900000001', 1200), (2, '华南', '13900000002', 300);" + "INSERT INTO feb VALUES (3, '华东', '13900000003', 900);") + + +async def test_a_union_of_monthly_tables_is_counted_with_the_sources_schema(monkeypatch, tmp_path): + """几张列一样的月度表用 SELECT * 连起来再汇总:按数据源的表结构数得出号码落在第几列,汇总的金额照常交给模型、照常 + 取值;选出全部列的,只遮号码那一列。* 展开数据源里的表又不按表结构数位置的话,展开落到的那一段整段都遮,金额也拿不到。""" + name, source_id = await _source(tmp_path, {"mask_columns": "buyer_phone"}, script=MONTHS) + try: + steps = [q(name, "SELECT region, SUM(amount) AS total FROM (SELECT * FROM jan UNION ALL SELECT * FROM feb) t " + "GROUP BY region ORDER BY total DESC", "c1"), + q(name, "SELECT * FROM jan UNION ALL SELECT * FROM feb ORDER BY id", "c2")] + script = Script(monkeypatch, steps, {"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "x"}) + out, _ = await ask(name) + assert out.error is None and out.result["value"] == 2100 + assert all(_clean(text) for text in script.tool_results().values()) + totals, rows = (artifact_store.load(query["artifact"]) for query in out.queries) + assert totals["masked_at_query"] == [] and totals["rows"] == [["华东", 2100], ["华南", 300]] + assert rows["masked_at_query"] == ["buyer_phone"] and _cells(rows, "amount") == [1200, 300, 900] + finally: + await engines.invalidate(source_id)