yilunjiang/AgentLab
1
0
Fork 0
mirror of https://github.com/ThearchyHelios/AgentLab.git synced 2026-11-15 03:41:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

合入期 5 前置修复 K3：核查 agent 的查询结果在工具边界先遮——取自遮罩列的结果列（别名、函数、子查询、WITH、*、UNION、表值函数、整行引用都认）在交给模型的内容和落库的快照里都换成「已遮罩」，模型看不到遮罩值；按值遮降为第二道防线；换了写法的长号码（空格、连字符、点、括号、全角、国家码）在回答、理由、工件和最后一道关里都认得出；认遮罩列的共用部分挪到 engine/masking.py

Browse Source
This commit is contained in:
ThearchyHelios
2026-10-06 22:19:42 +08:00
parent c1e1fcddfe b0fdf56eed
commit 58a6759fe2
15 changed files with 4805 additions and 1023 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/engine/inquiry.py
+210 -990
View File
File diff suppressed because it is too large Load Diff
backend/app/engine/masking.py
+2079
View File
File diff suppressed because it is too large Load Diff
backend/app/engine/nodes/challenge.py
+70 -2
View File
@@ -9,7 +9,11 @@
取自的查询算，合并查询按它合并的那几个源算）；这句没挂依据，就在这份报告引用的数据源里查。任务里给这句结论
和它挂的证据摘录（judge.prepare，数据源遮罩的列照旧不给；摘录里原报告的编号都标成「原报告 Qn」，和 agent 自己
的 Q1 分开），要它先想清楚「如果这句不成立，数据里会看到什么」，写只读查询去找，最后交结论（经受住 / 被削弱 /
被推翻 / 无法判断）和理由。理由里的数字只能引用它自己查询的格子，裸数字由 compose_doc 标出来
被推翻 / 无法判断）和理由。理由里的数字只能引用它自己查询的格子，裸数字由 compose_doc 标出来。
遮罩：agent 的查询结果在工具边界就遮好了；这份报告引用的查询快照里的遮罩值（按 SQL 认出改了名字的遮罩列，同审批
质询）反查之前就交给查询工具，agent 查到写着同一个号码的别的表时也看不到。摘录按列名遮，改了名字的遮罩列原值
agent 看得到，会照着写进理由：理由正文和理由文档按值换成「已遮罩」（换了写法的号码也认），这句结论原文里写着的
不换（_masked_rationale）
3. 出处：一条查询都没执行成功的，判什么都记为无法判断；判「被推翻」的，理由里还得引用到自己的查询结果——它要让
正式出具降档，审稿记录里得查得到凭什么。不采信时系统写明原因，反方原本的结论和理由照录在后面（纯文字：解析
不了的引用写成「（引用无效）」，裸数字注明「（无出处）」）。这句该查的数据源都不存在或已停用、核查模型本身用
@@ -44,8 +48,10 @@
"""
from __future__ import annotations
import asyncio
import re
from collections.abc import Callable, Collection, Iterable
from dataclasses import replace
from typing import Any
from langgraph.func import task
@@ -372,6 +378,54 @@ def _claim_of(cand: judging.Candidate, outcome: VerifyOutcome) -> tuple[dict[str
return _claim(cand, verdict, rationale, doc=outcome.rationale, queries=outcome.queries, usage=outcome.usage), None
def _spared(seed: Any, claim: str) -> Any:
"""报告快照里的遮罩值（seed）去掉这句结论原文里写着的，交给这一句的反查。seed 本身不动（下一句另算）。"""
known = seed.copy()
known.protect([claim])
return known
def _masked_rationale(outcome: VerifyOutcome, secrets: Any, *, given: Iterable[str | None]) -> VerifyOutcome:
"""反方的理由里的遮罩值换成「已遮罩」，换了写法的号码也认（199-1474-2785、+86 …）。值是这份报告引用的查询快照里的
（secrets，反查之前就交给了查询工具），加上 agent 查询时在工具边界遮掉的（outcome.masked_values）；被审的那句结论
原文里写着的（given）已经摆在读者眼前，不换。反方 agent 的查询结果在工具边界就遮好了，可证据摘录按列名遮，改了名字
的遮罩列（buyer_phone AS 手机号）的原值它看得到，会写进理由。
换的办法同审批质询的回答（_scrubbed_source）：原文里的值先换掉，再照原来的目录重新组装理由文档；引用到的格子按
遮好的快照渲染（文字格里夹着的遮罩值换掉那一截）。理由里认不出遮罩值的原样交回。"""
from app.engine.evidence import compose_doc
from app.engine.masking import _masked_view, _scrubbed_source, _Secrets
doc, result = outcome.rationale, outcome.result
# 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
pool.prepare()
source = str(doc.get("source") or "")
if not (pool.spans(str(doc.get("markdown") or ""), prose=True) or pool.spans(source, prose=True)):
return outcome
views: dict[str, Any] = {}
fetch = _fetcher()
for query in outcome.queries:
snapshot = fetch(query.get("artifact"))
if isinstance(snapshot, dict) and isinstance(snapshot.get("rows"), list):
views[str(query["artifact"])] = _masked_view(snapshot, set(), pool, embedded=True)[0]
def load(artifact: str) -> Any:
return views[artifact] if artifact in views else artifact_store.load(artifact)
try:
again = compose_doc(_scrubbed_source(source, pool), doc.get("catalog") or {},
node_id=str(doc.get("node_id") or ""), run_id=doc.get("run_id"), loader=load)
except Exception: # noqa: BLE001 - 重新组装出了意外，值也不能留着
# 正文按值换掉，理由文档不要了（界面照正文显示；判「被推翻」的因此查不到出处，按无法判断记）
said = pool.scrub(str(doc.get("markdown") or ""), prose=True)
return replace(outcome, rationale=None, result={**result, RATIONALE: said})
return replace(outcome, rationale=again, result={**result, RATIONALE: again["markdown"]})
def _cites(doc: dict[str, Any] | None) -> bool:
"""理由文档里有没有一处引用解析到了 agent 自己的查询结果（compose_doc 只把解析成功的引用记进 unit.cites）。"""
return isinstance(doc, dict) and any(unit.get("cites") for _, unit in iter_units(doc))
@@ -566,6 +620,9 @@ class _Review:
# 这句没挂依据时在整份报告引用的数据源里查：写作者看得到的就是这几个库
everywhere = claim_sources(self.catalog, self.catalog, self.fetch)
usable = await _usable_sources(everywhere)
# 报告引用的查询快照里遮罩列的值：反查之前交给查询工具（agent 查到写着同一个号码的别的表时看不到），反查之后
# 据此换掉理由里的遮罩值（_masked_rationale）
seed = await self._seed()
claims: list[dict[str, Any]] = []
usage = _zero_usage()
model: str | None = None
@@ -629,11 +686,14 @@ class _Review:
if allot.tightened:
self.run_seconds = allot.seconds
# 收紧过的：核查 agent 到点自己发的 verify_limit 改说运行时长上限将至（deadline.relabel），不写收紧后的秒数
# 这句结论原文里写着的值已经摆在读者眼前，不当遮罩值。上万个值过一遍要几十毫秒，放进线程
known = await asyncio.to_thread(_spared, seed, cand.text)
outcome = await run_verify_agent(
spec=self.spec, system=CHALLENGE_SYSTEM, task=task_text, sources=live,
output_schema=CHALLENGE_SCHEMA, max_queries=self.max_queries, timeout_s=allot.timeout_s,
daily_max_usd=self.daily_max_usd, emit=deadline.relabel(ctx.emit, allot), run_id=ctx.run.run_id,
node_id=ctx.node.id, data_versions=run_versions(ctx.run))
node_id=ctx.node.id, data_versions=run_versions(ctx.run), masked_values=known)
outcome = await asyncio.to_thread(_masked_rationale, outcome, known, given=[cand.text])
if outcome.error and not outcome.model:
broken = outcome.error # 连模型都没建成（没配好、接入出错）：换一句也一样
model = model or outcome.model
@@ -709,6 +769,14 @@ class _Review:
return {"artifact": artifact, "counts": counts, "usage": usage, "unreviewed": missed,
"claims": [{k: c[k] for k in ("unit", "text", "verdict")} for c in claims]}
async def _seed(self) -> Any:
"""这份报告目录里的查询快照里遮罩列的值（engine/masking 的 _Secrets）。收法同审批质询收报告快照
（inquiry.snapshot_secrets：按 SQL 认出改了名字的遮罩列，查询当时记下的遮罩加上数据源现在设的）。"""
from app.engine.inquiry import snapshot_secrets # inquiry 导入了这个模块，用到时再导入
return await snapshot_secrets(self.fetch(entry.get("artifact")) for entry in self.catalog.values()
if isinstance(entry, dict) and entry.get("kind") == "query")
async def _store(self, record: dict[str, Any]) -> str | None:
try:
return await artifact_store.put_json(record, kind="challenge", run_id=self.ctx.run.run_id,
backend/app/engine/nodes/dual_check.py
+11 -2
View File
@@ -13,7 +13,9 @@
连取值也不给；遮罩列不给。计算式只有一项来自查询、或者写了口径说明时不给列名；几项都来自查询又没写口径说明的，列名
里看不出日期、口径的取舍时才给，好让它分得清哪一项是哪一项。
**遮罩列上的输入不复算。** cell() 取自遮罩列的输入不能单独复算；别的输入（计算式里写成当时代进去的值）里按位置取自
**遮罩列上的输入不复算。** cell() 取自遮罩列的输入不能单独复算，改了名字、套了函数或汇总的遮罩列（SUM(amount) AS
gmv，amount 设了遮罩）按 SQL 认出来也一样（engine/masking.derived_columns）：盲算的查询在工具边界按同一套认法遮掉，
派出去也取不到数，不白花一次盲算。别的输入（计算式里写成当时代进去的值）里按位置取自
查询结果的一格（vars.res.rows[0][2]），顺着取值链认出它取自哪份快照的哪一列（recompute.input_origin，证据面板、推演、
可验证报告同一个认法），按需算那份快照要遮的列（和证据面板同一套，见 _Sources.hidden）——它不一定是 cell() 用到的
那几份。取自遮罩列、遮没遮核对不了的，整个指标按无法复算记：遮罩值不能当字面量写进交给盲算的计算式。伸进查询结果的行、
@@ -75,6 +77,7 @@ from app.engine.expressions import ExpressionError, _cell_call, _chain_root, cel
parse_expression
from app.engine.replay import _in_graph, ask, once
from app.engine.labels import DUAL_CHECK_FIELD_LABEL, DUAL_CHECK_ON_REJECT_LABEL
from app.engine.masking import derived_columns
from app.engine.schema import DUAL_CHECK_ON_REJECT, NodeType, dual_check_card
from app.engine.state import GraphState
@@ -613,9 +616,15 @@ async def _query_input(item: dict[str, Any], sources: _Sources, aliases: dict[st
value = cell_value(item.get("value"))
if not _number(value):
return {}, f"输入 {item['path']} 不是数值"
if column.lower() in await sources.hidden(artifact):
hidden = await sources.hidden(artifact)
if column.lower() in hidden:
return {}, f"输入 {item['path']} 取自设置了遮罩的列，不能用于复算"
found = await sources.source(source)
derived = derived_columns(snap, set(hidden), found.kind if found is not None else None) if hidden else None
if derived and column.lower() in derived:
# 改了名字、套了函数或汇总的遮罩列（SUM(amount) AS gmv）：盲算的查询在工具边界按同一套认法遮掉（engine/masking），
# 交回的格子注定取不到，不白派一次盲算。SQL 解析不了的认不准，照常派
return {}, f"输入 {item['path']} 由设置了遮罩的列算出（经改名、函数或汇总），盲算时会被遮罩，不能用于复算"
if found is None or not found.enabled:
return {}, f"数据源「{source}」不存在或已停用"
return {"path": item["path"], "key": f"{artifact}|{row}|{column}", "value": value,
backend/app/engine/verify_agent.py
+65 -12
View File
@@ -6,6 +6,16 @@
**只读。** 工具只建给定数据源的 db_schema / db_query，建的时候 force_readonly：可写源上的写语句也被守卫拒绝
（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。调用方交进来的和查询时遮掉的分开记（masking.SeenValues），合起来时 agent
查到的在前：合起来到了上限 SECRETS_MAX，挤掉的是报告里的值。
**数从快照取。** 每次成功的查询是一份查询快照（工具自己落的 put_json），按这个 agent 自己的顺序编 Q1、Q2…。
最后交结果时，调用方结构里的数值字段一律改成格子引用 {alias, row, column}，由代码按引用从快照里取值：模型自己
写的数字不采信，取不到就是 None，不兜底成 0。数值字段的各种写法都认（type 列表、anyOf 里的数、$ref 引用进来的、
@@ -45,7 +55,7 @@ import json
import logging
import time
from collections.abc import Callable, Iterable
from dataclasses import asdict, dataclass, field
from dataclasses import asdict, dataclass, field, replace
from enum import Enum
from typing import Any
@@ -59,13 +69,14 @@ 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 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
from app.providers import catalog as pricing
from app.providers.budget import BudgetFallback, budget_refused, cut_off, cut_words, error_usage, truncated_error
from app.providers.factory import ModelSpec, ProviderNotConfigured, bind_tools_safely, get_chat_model
from app.tools.datasource import QUERY_PREFIX, SCHEMA_PREFIX, build_datasource_tools
from app.tools.datasource import MASKED_VALUES_KEY, QUERY_PREFIX, SCHEMA_PREFIX, build_datasource_tools
from app.tools.registry import ToolArgsError, ToolContext, Versions, args_model_of, prepare_args
logger = logging.getLogger(__name__)
@@ -103,6 +114,8 @@ VERIFY_SYSTEM = (
"- 先用 db_schema 工具看表结构，确认表名、字段名，再用 db_query 工具写只读 SQL。不能改数据，写语句会被拒绝\n"
"- 每次成功的查询前面都标着编号【证据 Qn】，最后交结果时用它说明每个数取自哪次查询的哪一格\n"
"- {limit}先想清楚要查什么，不要重复同样的查询\n"
"- 数据源设了遮罩的列（手机号、姓名这类），查询结果里显示为「已遮罩」：不要设法绕过，交结果、写理由时也不要引用"
"这些格子\n"
"- 不要编造数字，查不到就照实说"
)
VERIFY_FINAL_PROMPT = (
@@ -149,6 +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），连同调用方交进来的：agent 查到的在前，合起来
到了上限 SECRETS_MAX 挤掉的是调用方的（SeenValues.pooled）。只在内存里交给调用方按值遮（审批质询的回答、反方审稿的
理由），不进 as_dict、不往外交。
"""
result: dict[str, Any] | None
@@ -162,9 +178,12 @@ class VerifyOutcome:
answered: bool = False
notes: list[str] = field(default_factory=list)
duration_ms: int = 0
masked_values: Any = field(default=None, repr=False, compare=False)
def as_dict(self) -> dict[str, Any]:
return asdict(self)
out = asdict(replace(self, masked_values=None))
out.pop("masked_values", None)
return out
# --------------------------------------------------------------------------
@@ -407,6 +426,18 @@ class _Cells:
self.cells.append({"path": path, "status": "unresolved", "alias": alias, "artifact": query["artifact"],
"reason": f"{alias} {e}"})
return None
if str(name).lower() in _masked_at_query(content):
# 查询时就遮了的列：格子里是「已遮罩」，不是数据源里的值。照实说是遮罩列，不说成「不是数值」
self.cells.append({"path": path, "status": "unresolved", "ref": f"{alias}.r{row}.{name}", "alias": alias,
"row": row, "column": name, "artifact": query["artifact"],
"reason": f"{alias} 的「{name}」取自设置了遮罩的列，查询时已遮罩，不采信"})
return None
if raw == _masked_word() and isinstance(content, dict) and "masked_at_query" in content:
# 这一格查询时按值换掉了（和别的查询、报告里的遮罩值相同）：同样照实说
self.cells.append({"path": path, "status": "unresolved", "ref": f"{alias}.r{row}.{name}", "alias": alias,
"row": row, "column": name, "artifact": query["artifact"],
"reason": f"{alias} 第 {row} 行「{name}」与设置了遮罩的值相同，查询时已遮罩，不采信"})
return None
truth = _typed(raw, {"type": sorted(types)}, column_kind(content, name))
cited = {"ref": f"{alias}.r{row}.{name}", "alias": alias, "row": row, "column": name,
"artifact": query["artifact"], "eid": cell_eid(query["artifact"], row, name), "value": truth}
@@ -502,6 +533,12 @@ class _Cells:
return None
def _masked_at_query(content: Any) -> set[str]:
"""查询快照里查询时就整列遮掉的列（小写，tools/datasource 记下的 masked_at_query）。"""
listed = content.get("masked_at_query") if isinstance(content, dict) else None
return {c.lower() for c in listed if isinstance(c, str)} if isinstance(listed, list) else set()
# --------------------------------------------------------------------------
# 花费
# --------------------------------------------------------------------------
@@ -636,6 +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 并进来）。后面的
#: 查询按它们按值遮，最后合起来交回调用方（run_verify_agent），agent 查到的在前
seen: SeenValues = field(default_factory=SeenValues)
def outcome(self, **fields: Any) -> VerifyOutcome:
usage = self.usage
@@ -665,6 +705,7 @@ async def run_verify_agent(
run_id: str = "",
node_id: str = "",
data_versions: dict[str, Any] | Versions | None = Versions.UNSET,
masked_values: Any = None,
) -> VerifyOutcome:
"""让一个只读 agent 查库，按 output_schema 交回有出处的结果。永远不抛。
@@ -679,16 +720,27 @@ async def run_verify_agent(
tool.end（role: verify，查询带 query_artifact）；不在节点里（质询接口）就不传，什么都不发。
- run_id / node_id：查询快照、表结构快照记在谁名下。data_versions：上传表格用哪一版（运行固定的那份
data_versions；不在运行里传 Versions.CURRENT；缺省 UNSET 时上传源一律报错，不拿当前版本顶上）。
- masked_values：调用方已经认出的遮罩值（engine/masking 的 _Secrets，审批质询、反方审稿收的报告快照里的值）。
查询工具从第一条查询起就按它们按值遮：agent 查到的备注、别的表里写着同一个号码时，模型看不到。随查询时遮掉的值
一起交回（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:
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],
@@ -728,12 +780,13 @@ async def _verify(state: _State, *, spec: ModelSpec, system: str, task: str, sou
notes.append(note)
await _announce(emit, EventType.LOG, level="warn", code="verify_unpriced", message=note)
# ---- 工具：只读的 db_schema / db_query
# ---- 工具：只读的 db_schema / db_query，查询结果先遮再交给模型（遮掉的值并进 state.seen）
wanted = list(dict.fromkeys(s for s in sources if isinstance(s, str) and s))
ctx = ToolContext(run_id=run_id, node_id=node_id, data_versions=data_versions, emit=emit)
ctx = ToolContext(run_id=run_id, node_id=node_id, data_versions=data_versions, emit=emit,
extra={MASKED_VALUES_KEY: state.seen})
async with SessionLocal() as session:
tools = await build_datasource_tools([f"{p}{s}" for s in wanted for p in (SCHEMA_PREFIX, QUERY_PREFIX)], ctx,
session, force_readonly=True)
session, force_readonly=True, mask_results=True)
tool_map = {t.name: t for t in tools}
if lost := [s for s in wanted if f"{QUERY_PREFIX}{s}" not in tool_map]:
notes.append("数据源" + "、".join(f"「{s}」" for s in lost) + "不存在或已停用，核查没有查它")
backend/app/tools/datasource.py
+50 -7
View File
@@ -53,6 +53,10 @@ RUN_VERSIONS_KEY = "data_versions"
DATA_UNAVAILABLE = "查询失败（数据不可用）："
#: 运行中途补固定一个源时发的那条日志事件的 code（_pin_late）
LATE_PIN_CODE = "data_version_pinned"
#: ToolContext.extra 里放「认出的遮罩值」的键（engine/masking 的 SeenValues：调用方交进来的，加上查询时遮掉的）：核查
#: agent（engine/verify_agent.py）放进去，打开了遮罩开关的查询工具按它们按值遮，再把每次遮掉的值并进去。只在内存里给
#: 调用方按值遮用，不进快照、不交给模型
MASKED_VALUES_KEY = "masked_values"
#: 每次查询之后基于数据目录的 SQL 检查最多等这么久。正常一条 SQL 几毫秒；等不到就放弃这一次检查，
#: 查询结果照常交回——检查是附加的，不能拖慢查询
@@ -199,7 +203,8 @@ def _query_description(source: Any, *, readonly: bool | None = None) -> str:
async def build_datasource_tools(
names: list[str], ctx: ToolContext, session: AsyncSession, *, force_readonly: bool = False
names: list[str], ctx: ToolContext, session: AsyncSession, *, force_readonly: bool = False,
mask_results: bool = False,
) -> list[StructuredTool]:
"""按工具名反查数据源并构造工具。build_tools 的第四个来源。
@@ -209,6 +214,9 @@ async def build_datasource_tools(
force_readonly：查询工具一律按只读守卫检查，可写源上的写语句也拒绝（见 _make_query_tool）。
核查（engine/verify_agent.py）建工具时传它：那一路没有审批关卡，绝不能写。
mask_results：查询结果交给模型、落进快照之前先遮（见 _make_query_tool）。核查建工具时也传它：模型看不到遮罩列的
值，就不会把它们写进回答、理由。缺省关着：Agent 节点、调用工具节点的查询照旧原样。
"""
wanted: dict[str, list[str]] = {}
for name in names:
@@ -238,7 +246,8 @@ async def build_datasource_tools(
entries = await _catalog_of(row)
for prefix in prefixes:
if prefix == QUERY_PREFIX:
tools.append(_make_query_tool(source, ctx, fixed=fixed, catalog=entries, force_readonly=force_readonly))
tools.append(_make_query_tool(source, ctx, fixed=fixed, catalog=entries, force_readonly=force_readonly,
mask_results=mask_results))
else:
tools.append(_make_schema_tool(source, fixed=fixed, catalog=entries))
return tools
@@ -450,7 +459,7 @@ def _unavailable_tool(row: DataSource, prefix: str, reason: str) -> StructuredTo
def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
catalog: Mapping[str, data_catalog.CatalogEntry] | None = None,
force_readonly: bool = False) -> StructuredTool:
force_readonly: bool = False, mask_results: bool = False) -> StructuredTool:
"""source 是手工源的 DataSource，或上传源绑定快照的 SourceView（字段同名）。
fixed：上传源的版本是这次运行固定的（找不到时的说法不同）。
@@ -458,6 +467,10 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
force_readonly：不管源可不可写，一律按只读守卫检查，写语句交回 FORCED_READONLY、不执行。只靠守卫这一道：
连接池按源缓存、不分只读与否（engine.cache_key），可写源借不到连接层只读的连接——别为此另起一个只读池，
缓存键相同，会把这个源平时用的可写连接也换成只读的。
mask_results：数据源设了遮罩时，取自遮罩列的格子（按 SQL 认，改了名字的也算）在交给模型、落进快照之前就换成
「已遮罩」，见 engine/masking.py 的 mask_query_result；这一次遮掉的值并进 ctx.extra[MASKED_VALUES_KEY]（有的话）。
那里已经有的值（这个 agent 前面的查询遮掉的、调用方交进来的）这一次也按值换，数据源没设遮罩也一样：先查号码、
再查写着号码的备注或别的库，后一份同样遮好。查询失败时报错原话里的值也换（masking.mask_error_text）。
"""
# 数据源在 options.query_timeout_s 里配了就用它，否则取缺省。数据库按它停下语句，
# 引擎按它告诉界面上限是多少（metadata["timeout_s"]）
@@ -485,6 +498,12 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
if is_write(statement, dialect=(source.kind or "").lower()):
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 = 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:
@@ -501,10 +520,29 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
# 驱动的原话留着（"no such table: x" 正是模型改写 SQL 要的线索），
# 类名和 SQLAlchemy 的包装前缀去掉：这一行也会原样出现在运行面板上
return f"查询失败：{first_line(e) or explain(e)[0]}"
said = first_line(e) or explain(e)[0]
if mask_results and (mask or known is not None):
from app.engine.masking import mask_error_text
# 数据库会把转换失败的值写进报错（invalid input syntax for type bigint: "139-0000-0001"）
said = mask_error_text(said, sql, known)
return f"查询失败：{said}"
payload = result.to_payload()
payload["source"] = source.name
# 查询时就整列遮掉的结果列；没开遮罩开关、数据源没设遮罩又没有可按值换的时为 None（快照和原来一字不差）
masked_now: list[str] | None = None
if mask_results and (mask or known is not None):
from app.engine.masking import mask_query_result
# 按 SQL 认列、逐格比对是纯计算，放进线程（同 _sql_checks）
done = await asyncio.to_thread(mask_query_result, payload, mask, source.kind, known,
schema=source.schema_cache)
if mask or done.payload is not payload:
payload, masked_now = done.payload, done.columns
if sink is not None:
# 并进 agent 查到的那一份（不和调用方交进来的混在一起，上限收满了也不丢），也放进线程
await asyncio.to_thread(sink.add, done.secrets)
# 对照数据目录检查这条 SQL（一对多关联后重复计算、存量跨期求和……）。结果交给模型（Agent 据此改写），
# 也随查询快照存下：证据面板看得到，口径卡据此给指标标「存疑」、出具降档。没查出问题不加这个键，
# 快照和以前一字不差
@@ -518,14 +556,19 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
try:
from app.core.artifact_store import put_json
# 只进快照、不进交给模型的结果的两样（模型看到多余的字段只会误会：以为数据被遮了，
# 只进快照、不进交给模型的结果的几样（模型看到多余的字段只会误会：以为数据被遮了，
# 或者把一串十六进制抄进报告）：
# - 查询当时数据源设的遮罩列：数据源事后改名、删掉，证据面板仍按当时的遮；
# - 查询当时数据源设的遮罩列：数据源事后改名、删掉，证据面板仍按当时的遮。只记数据源设的列名——证据面板把
# 同源各份快照记的这一项并起来，从表的字段清单里去掉，混进结果列的别名会把同名的真实字段也去掉；
# - masked_at_query：开了遮罩开关时，查询时就整列遮掉的结果列（这几列的格子是「已遮罩」，不是数据源里的值，
# 改了名字的也在这里）。有这一项的快照，别的格子整格是「已遮罩」的是查询时按值换掉的；
# - data_version：上传表格查的是哪个快照，事后据此认得出这个数出自哪一版数据。
# 手工源、没设遮罩的，快照和原来一字不差
extra: dict[str, Any] = {}
if mask := masked_columns(source.options):
if mask:
extra["mask_columns"] = mask
if masked_now is not None:
extra["masked_at_query"] = masked_now
if version := getattr(source, "snapshot_id", None):
extra["data_version"] = version
# 工件引用的 meta 记下源和 SQL 里用到的表：数据目录按它数每张表被运行查询过几次
backend/tests/test_challenge_node.py
+131
View File
@@ -1102,6 +1102,114 @@ async def test_the_task_gives_the_claim_its_evidence_and_only_its_sources(monkey
await engines.invalidate(source_id)
async def test_the_challenge_agent_never_sees_masked_values(monkeypatch, tmp_path):
"""反查 agent 自己的查询选出了遮罩列（起了别名）：交给它的工具结果、它的查询快照里都是「已遮罩」，理由里引用到那一格
渲染出来也是「已遮罩」；审稿记录、运行事件里都没有原值。以前查询结果原样交给模型、原样存下，理由也原样渲染。"""
name, source_id = await _source(tmp_path, options={"mask_columns": ["buyer_phone"]})
try:
top = "SELECT buyer_phone AS 电话, amount AS 金额 FROM sales WHERE region = '华东' ORDER BY amount DESC"
plays = {"居各区域之首": Play([top], "stands", "华东最大一笔来自 [[v:Q1.r0.电话]]，金额 [[v:Q1.r0.金额]]。[[see:Q1]]")}
models = Models(monkeypatch, name, plays=plays)
row = await finish(graph(name), run_class="exploratory")
assert row.status == "succeeded", row.error
seen = [str(m.content) for _, key, messages in models.calls if key == "居各区域之首"
for m in messages if isinstance(m, ToolMessage)]
assert seen and any("已遮罩" in s for s in seen) and all("1390000000" not in s for s in seen)
record = artifact_store.load((await output_of(row.id, "chal"))["challenge_artifact"])
[claim] = [c for c in record["claims"] if "居各区域之首" in c["text"]]
assert claim["verdict"] == "stands" and claim["rationale"] == "华东最大一笔来自 已遮罩，金额 1,200。"
snapshot = artifact_store.load(claim["queries"][0]["artifact"])
assert {r[0] for r in snapshot["rows"]} == {"已遮罩"} and snapshot["masked_at_query"] == ["电话"]
assert "1390000000" not in json.dumps(record, ensure_ascii=False)
async with SessionLocal() as session:
evs = (await session.execute(select(RunEvent).where(RunEvent.run_id == row.id))).scalars().all()
assert evs and all("1390000000" not in json.dumps(e.data, ensure_ascii=False) for e in evs)
finally:
await engines.invalidate(source_id)
async def test_masked_values_written_into_the_rationale_are_masked_in_any_shape(monkeypatch, tmp_path):
"""复核第 1 轮：报告取数把遮罩列改了名字选出来（buyer_phone AS 手机号），摘录按列名遮不住，反方 agent 看得到号码，
把它换个写法写进理由（139-0000-0001、+86 139 0000 0002）。审稿记录里的理由正文、理由文档都换成「已遮罩」：值从这句
挂的报告快照里按 SQL 认出来收（和审批质询收报告快照的办法一样），连同 agent 查询时在工具边界遮掉的值；同时报告
快照里的值在反查之前就交给查询工具，agent 查到写着同一个号码的别的表时也看不到。以前理由原样存下、随运行封存。"""
name, source_id = await _source(tmp_path, options={"mask_columns": ["buyer_phone"]})
try:
sql = "SELECT region AS 区域, buyer_phone AS 手机号, amount AS 金额 FROM sales ORDER BY id"
said = "来自 139-0000-0001 号码的那笔属实，另有 +86 139 0000 0002 的一笔。"
plays = {"华东有一笔": Play(["SELECT COUNT(*) AS 笔数 FROM sales WHERE region = '华东'"], "stands", said)}
Models(monkeypatch, name, writer=("华东有一笔 [[v:Q1.r0.金额]] 的订单。[[see:Q1]]",), plays=plays)
row = await finish(graph(name, sql=sql), run_class="exploratory")
assert row.status == "succeeded", row.error
record = artifact_store.load((await output_of(row.id, "chal"))["challenge_artifact"])
[claim] = [c for c in record["claims"] if "华东有一笔" in c["text"]]
assert claim["verdict"] == "stands"
assert claim["rationale"] == "来自 已遮罩 号码的那笔属实，另有 已遮罩 的一笔。"
assert claim["rationale_doc"]["markdown"] == claim["rationale"]
assert claim["rationale_doc"]["source"] == claim["rationale"], "存下的原文也换掉"
shown = json.dumps(record, ensure_ascii=False)
assert "0000-0001" not in shown and "0000 0002" not in shown
finally:
await engines.invalidate(source_id)
async def test_the_challenge_agent_does_not_see_report_values_in_other_tables(monkeypatch, tmp_path):
"""报告取数选了手机号（遮罩列）；反方 agent 去查回访记录，备注里写着同一个号码（备注没设遮罩）。报告快照里的遮罩值
在反查之前就交给了查询工具：交给模型的备注里那一截已经是「已遮罩」，它的查询快照里也是。以前原样交给模型。"""
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
db = sqlite3.connect(path)
db.executescript(
"CREATE TABLE sales (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, amount INTEGER, refunded INTEGER);"
"INSERT INTO sales VALUES (1, '华东', '13900000001', 1200, 0), (2, '华南', '13900000003', 900, 0);"
"CREATE TABLE callbacks (id INTEGER PRIMARY KEY, region TEXT, note TEXT);"
"INSERT INTO callbacks VALUES (1, '华东', '客户 139-0000-0001 来电改址'), (2, '华南', '例行回访');")
db.commit()
db.close()
name = f"shop_{uuid.uuid4().hex[:8]}"
async with SessionLocal() as session:
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=True, enabled=True,
options={"mask_columns": ["buyer_phone"]})
session.add(row)
await session.flush()
row.schema_cache = await introspect.introspect(row)
await session.commit()
source_id = row.id
try:
sql = "SELECT region AS 区域, buyer_phone, amount AS 金额 FROM sales ORDER BY id"
notes = "SELECT region AS 区域, note AS 备注 FROM callbacks ORDER BY id"
plays = {"华东有一笔": Play([notes], "stands", "回访记录里有 [[v:Q1.r0.备注]]。[[see:Q1]]")}
models = Models(monkeypatch, name, writer=("华东有一笔 [[v:Q1.r0.金额]] 的订单。[[see:Q1]]",), plays=plays)
row = await finish(graph(name, sql=sql), run_class="exploratory")
assert row.status == "succeeded", row.error
seen = [str(m.content) for _, key, messages in models.calls if key == "华东有一笔"
for m in messages if isinstance(m, ToolMessage)]
assert any("客户 已遮罩 来电改址" in s for s in seen) and all("0000-0001" not in s for s in seen)
record = artifact_store.load((await output_of(row.id, "chal"))["challenge_artifact"])
[claim] = [c for c in record["claims"] if "华东有一笔" in c["text"]]
assert claim["rationale"] == "回访记录里有 客户 已遮罩 来电改址。"
snapshot = artifact_store.load(claim["queries"][0]["artifact"])
assert snapshot["rows"] == [["华东", "客户 已遮罩 来电改址"], ["华南", "例行回访"]]
finally:
await engines.invalidate(source_id)
def test_values_written_in_the_claim_itself_are_kept_in_the_rationale():
"""被审的那句结论原文里写着的号码已经摆在读者眼前：理由里复述它不换，别的号码照换。"""
from app.engine.evidence import compose_doc
from app.engine.masking import _Secrets
from app.engine.verify_agent import VerifyOutcome
secrets = _Secrets()
secrets.add({"columns": ["p"], "rows": [["13900000001"], ["13900000003"]]}, {"p"})
doc = compose_doc("号码 13900000003 的客户在华南，139-0000-0001 不在。", {}, node_id="chal", run_id="r1")
outcome = VerifyOutcome(result={"verdict": "stands", "rationale": doc["markdown"]}, cells=[], queries=[],
rationale=doc, usage={})
masked = challenge_mod._masked_rationale(outcome, secrets, given=["号码 13900000003 的客户在华南。"])
assert masked.result["rationale"] == "号码 13900000003 的客户在华南，已遮罩 不在。"
assert masked.rationale["markdown"] == masked.result["rationale"] and masked.rationale["node_id"] == "chal"
assert challenge_mod._masked_rationale(outcome, _Secrets(), given=[]) is outcome, "没有遮罩值时原样"
async def test_a_metric_excerpt_marks_the_report_query_numbers_too(monkeypatch, shop):
"""句子挂的是口径卡指标（正式报告的主路径「取数 → 口径卡 → 报告」）：摘录正文里「取自查询 Q1」「查询 Q1 的 SQL」
也标成原报告的编号，和反查 agent 自己的 Q1 分开；工具按指标的输入取自的查询找源。"""
@@ -1297,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"] == "号码 已遮罩 的客户在华南。"
backend/tests/test_dual_check_node.py
+53
View File
@@ -1034,6 +1034,59 @@ async def test_a_blind_cell_on_a_masked_column_is_not_used(monkeypatch, masked_s
assert "30.25" not in json.dumps(out, ensure_ascii=False)
async def test_a_renamed_masked_column_in_the_blind_query_is_masked_before_the_blind_agent_sees_it(
monkeypatch, masked_shop):
"""盲算的查询选出了改了名字的遮罩列（region AS 区域）：交给盲算模型、落进盲算快照之前就是「已遮罩」，快照记下查询
时遮了哪一列；没碰遮罩列的那一列（销售额）照常取值、照常比对。"""
name, _ = masked_shop
east = {**GMV, "id": "east", "name": "华东销售额", "expression": "cell(nodes.fetch, 0, 'gmv')"}
renamed = (f"SELECT region AS 区域, SUM(amount) AS gmv FROM orders WHERE status <> 'refunded' AND paid_date {WEEK} "
"GROUP BY region ORDER BY region")
Blind(monkeypatch, name, {"华东销售额": (renamed, "gmv")})
run = await start(graph(name, BY_REGION, metrics=[east], dual={"definitions": {}}))
assert run.status == "succeeded", run.error
evs = await events(run.id)
out = dual_output(run, evs)
assert out["status"] == "agreed" and out["metrics"][0]["status"] == "agreed"
[end] = [e.data for e in evs if e.type == "tool.end" and e.node_id == "dual" and e.data.get("query_artifact")]
snapshot = artifact_store.load(end["query_artifact"])
assert {row[0] for row in snapshot["rows"]} == {"已遮罩"} and snapshot["masked_at_query"] == ["区域"]
shown = json.loads(end["preview"].split("\n", 1)[1])
assert {row[0] for row in shown["rows"]} == {"已遮罩"}, "交给盲算模型的结果里也是「已遮罩」"
async def test_a_metric_summed_over_a_masked_column_is_not_blind_computed(monkeypatch, tmp_path):
"""复核第 1 轮：金额这一列设了遮罩，指标是 SUM(amount)。核查 agent 的查询结果在工具边界遮好，汇总遮罩列的值也遮
（只有一行的汇总就是那一行的原值），盲算交回的格子注定取不到。所以派出盲算之前按 SQL 认出这一列由遮罩列算出，
直接按无法复算记、说清原因，不白花一次盲算的钱。以前只按列名认，gmv 认不出，盲算照样派出去，最后才在格子引用上
被拒，原因说成「盲算交回的格子引用不可用」。"""
name, source_id = await add_source(build(tmp_path / "amount_masked.db"), options={"mask_columns": "amount"})
try:
blind = Blind(monkeypatch, name, {"销售额": (BLIND_TOTAL, "blind_total")})
run = await start(graph(name, RIGHT))
evs = await events(run.id)
out = dual_output(run, evs)
[row] = out["metrics"]
assert row["status"] == "unchecked" and blind.runs == 0, "没派盲算"
assert row["reason"] == ("输入 cell(nodes.fetch, 0, 'gmv') 由设置了遮罩的列算出（经改名、函数或汇总），"
"盲算时会被遮罩，不能用于复算")
assert not [e for e in evs if e.type == "tool.end" and e.node_id == "dual"]
finally:
await drop_source(source_id)
async def test_a_metric_filtered_by_a_masked_column_inside_its_sum_is_still_blind_computed(monkeypatch, masked_shop):
"""遮罩列（区域）只在汇总的 CASE 条件里（SUM(CASE WHEN region = '华东' THEN amount ELSE 0 END)）：同写在 WHERE 里，
值全部来自金额，原 SQL 和盲算的查询都不遮，照常比对。"""
name, _ = masked_shop
east = (f"SELECT SUM(CASE WHEN region = '华东' THEN amount ELSE 0 END) AS gmv FROM orders "
f"WHERE status <> 'refunded' AND paid_date {WEEK}")
blind = Blind(monkeypatch, name, {"销售额": (east, "gmv")})
run = await start(graph(name, east))
out = dual_output(run, await events(run.id))
assert blind.runs == 1 and out["metrics"][0]["status"] == "agreed", out["metrics"][0]
async def test_masked_group_columns_are_not_shown_to_the_blind_agent(monkeypatch, masked_shop):
name, _ = masked_shop
first = {**GMV, "id": "first", "name": "首组销售额", "expression": "cell(nodes.fetch, 0, 'gmv')"}
backend/tests/test_inquiry_api.py
+332 -10
View File
@@ -904,6 +904,78 @@ async def test_resuming_after_an_inquiry_does_not_repeat_the_nodes(monkeypatch,
ALIASED = "SELECT buyer_phone AS 手机号, amount AS 金额 FROM sales ORDER BY amount DESC LIMIT 1"
async def test_the_inquiry_agent_sees_and_stores_its_queries_masked(monkeypatch, masked_shop, client):
"""质询 agent 的查询结果在交给模型之前就遮好了（核查 agent 建查询工具时打开遮罩开关）：模型看到的工具结果、落库的
查询快照、运行详情「工件」页签走的通用工件接口里，遮罩列（起了别名的也算）都是「已遮罩」。以前快照原样存着，模型
看得到原值，通用接口也原样给出。回答照常按值遮（查询时遮掉的值只在内存里交回来）。"""
plays = {"买得最多": Play([ALIASED], "这一笔来自 [[v:Q1.r0.手机号]]（13900000001），金额 [[v:Q1.r0.金额]]。[[see:Q1]]")}
Models(monkeypatch, masked_shop, 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)
paused = await pause(masked_shop)
[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
assert r.json()["inquiry"]["answer"]["text"] == "这一笔来自 已遮罩（已遮罩），金额 1,200。"
[query] = r.json()["inquiry"]["queries"]
assert seen and any("已遮罩" in s for s in seen) and all("1390000000" not in s for s in seen)
got = await client.get(f"/api/artifacts/{query['artifact']}")
assert got.status_code == 200 and "1390000000" not in got.text
content = got.json()["content"]
assert content["rows"] == [["已遮罩", 1200]] and content["masked_at_query"] == ["手机号"]
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。"""
@@ -1168,6 +1240,48 @@ async def test_a_masked_column_after_a_label_in_a_union_keeps_its_values(monkeyp
await engines.invalidate(source_id)
_BIG_ONE = "CASE WHEN amount > 1000 THEN 1 ELSE 0 END"
@pytest.mark.parametrize("shadowed", [
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT *, {_BIG_ONE} AS vip_level FROM sales) t", id="* 在前"),
pytest.param(f"SELECT x.vip_level AS 档 FROM (SELECT * FROM sales s JOIN (SELECT id AS bid, {_BIG_ONE} AS vip_level "
"FROM sales) b ON s.id = b.bid) x", id="JOIN 子查询 * 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT (vip_level), {_BIG_ONE} AS vip_level FROM sales) t",
id="(col) 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT vip_level COLLATE NOCASE, {_BIG_ONE} AS vip_level FROM sales) t",
id="COLLATE 在前"),
])
async def test_a_masked_column_shadowing_a_label_keeps_its_values(monkeypatch, tmp_path, client, shadowed):
"""第 3 轮复验（J2 留下的缺口）：派生来源里 * 带出的等级、套了括号或 COLLATE 的等级排在同名的标签前面，SQLite 交回
的是真实等级。以前这一列被当成只装着标签，等级 1 撞上标签 1 不收，之后 WHERE vip_level = 1 在查询清单、事件、质询
工件里都照原样。现在不算只装着标签的列，两条都遮。"""
name, source_id = await _source(tmp_path, options={"mask_columns": "vip_level"}, script=LEVELS)
first, second = (f"SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = {level}" for level in (1, 2))
plays = {"各等级": Play([shadowed, first, second],
"问到的两个等级各有 [[v:Q2.r0.人数]] 人、[[v:Q3.r0.人数]] 人。[[see:Q2]]")}
Models(monkeypatch, name, writer="华东首单金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
try:
paused = await pause(name, graph=fetching(name, AMOUNTS))
[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, "问到的两个等级各有 2 人、1 人。")
sealed = [shadowed, *["SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = 已遮罩"] * 2]
assert [q["sql"] for q in item["queries"]] == sealed
[event] = await events(paused.id, "approval.inquiry")
assert [q["sql"] for q in event.data["queries"]] == sealed
for query, sql in zip(item["queries"], sealed):
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
assert got.status_code == 200 and got.json()["content"]["sql"] == sql, "质询工件的 SQL 和事件里的一字不差"
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{item['queries'][0]['artifact']}")
assert {cell for [cell] in got.json()["content"]["rows"]} == {"已遮罩"} and got.json()["masked"] == ["档"]
finally:
await engines.invalidate(source_id)
STAR_TAIL = "SELECT *, SUBSTR(buyer_phone, 4) FROM sales ORDER BY amount DESC LIMIT 1"
@@ -1403,14 +1517,22 @@ CALLBACKS = SALES + ("CREATE TABLE callbacks (id INTEGER PRIMARY KEY, region TEX
LEAKED = "回答中含有设了遮罩的值，无法遮住，已不予展示，请换个问法重新提问"
async def test_a_masked_value_inside_a_text_cell_is_cut_out_of_the_artifact(monkeypatch, tmp_path, client):
"""备注里写着手机号（备注没设遮罩，手机号设了遮罩、被报告取数选中）：agent 引用这一格，回答照旧整份不交；
质询工件里这份快照照样记着、点得开，那一格里的号码换成「已遮罩」，别的字、别的格子照原样，不然「不予展示」
旁边一点就能看到原值。"""
async def test_a_masked_value_inside_a_text_cell_is_cut_out_before_the_agent_sees_it(monkeypatch, tmp_path, client):
"""备注里写着手机号（备注没设遮罩，手机号设了遮罩、被报告取数选中）：报告快照里的遮罩值在开始回答之前就交给核查
agent 的查询工具，agent 查备注时那一截就换成「已遮罩」，模型看不到；它引用这一格，回答照常给出，渲染出来也是遮好的。
质询工件、通用工件接口里这份快照同样遮好，别的字、别的格子照原样。以前模型看得到原值，回答整份不交。"""
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CALLBACKS)
notes = "SELECT region AS 区域, note AS 备注 FROM callbacks ORDER BY id"
plays = {"回访": Play([notes], "华东的回访备注是 [[v:Q1.r0.备注]]。[[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:
report_sql = "SELECT buyer_phone, amount AS 金额 FROM sales ORDER BY amount DESC"
paused = await pause(name, graph=fetching(name, report_sql))
@@ -1420,20 +1542,148 @@ async def test_a_masked_value_inside_a_text_cell_is_cut_out_of_the_artifact(monk
assert r.status_code == 200, r.text
assert "1390000000" not in r.text
item = r.json()["inquiry"]
assert (item["answer"]["ok"], item["answer"]["text"], item["answer"]["doc_artifact"]) == (False, "", None)
assert item["answer"]["error"] == LEAKED
assert (item["answer"]["ok"], item["answer"]["text"]) == (True, "华东的回访备注是 客户已遮罩来电改址。")
[query] = item["queries"]
assert query["sql"] == notes
[event] = await events(paused.id, "approval.inquiry")
assert event.data["inquiry_artifacts"] == [query["artifact"]]
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
assert got.status_code == 200, got.text
assert got.json()["content"]["rows"] == [["华东", "客户已遮罩来电改址"], ["华南", "例行回访"]]
assert got.json()["masked"] == [] and "1390000000" not in got.text
assert any("客户已遮罩来电改址" in s for s in seen) and all("1390000000" not in s for s in seen), \
"交给模型的查询结果里就是遮好的"
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, got.text
assert got.json()["content"]["rows"] == [["华东", "客户已遮罩来电改址"], ["华南", "例行回访"]]
assert "1390000000" not in got.text
assert got.json().get("masked", []) == [], "备注只换了一截，不画成遮罩列"
finally:
await engines.invalidate(source_id)
#: 质询 agent 的两条查询：一条选出手机号（查询时遮掉），一条查回访备注（备注没设遮罩，里面写着手机号）
PHONES_EAST = "SELECT buyer_phone AS 号码, amount AS 金额 FROM sales WHERE region = '华东' ORDER BY amount DESC"
NOTES = "SELECT region AS 区域, note AS 备注 FROM callbacks ORDER BY id"
async def test_a_number_from_the_agents_other_query_is_masked_when_the_number_came_first(monkeypatch, tmp_path, client):
"""复核第 1 轮（对基线回退）：报告取数没选手机号，号码只出现在 agent 自己的另一次查询里。先查号码、后查备注：号码
查询时就遮掉、值并进内存，后面那次查询的备注在交给模型、落库之前就按值换掉，回答照常给出；质询工件、通用工件接口
给的都是遮好的。以前事后取质询工件时收不到这个号码，备注原样给出。"""
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CALLBACKS)
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, 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, "华东的回访备注是 客户已遮罩来电改址。")
phones, notes = item["queries"]
for url in (f"/api/runs/{paused.id}/inquiries/artifacts/{notes['artifact']}",
f"/api/artifacts/{notes['artifact']}"):
got = await client.get(url)
assert got.status_code == 200, got.text
assert got.json()["content"]["rows"] == [["华东", "客户已遮罩来电改址"], ["华南", "例行回访"]]
assert "1390000000" not in got.text
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{phones['artifact']}")
assert got.json()["content"]["rows"] == [["已遮罩", 1200], ["已遮罩", 300]] and got.json()["masked"] == ["号码"]
[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_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):
"""先查备注、后查号码：查备注时还不知道里面写着遮罩值，原样交给模型、原样存下；回答引用这一格，渲染出来才露，整份
不交。回答时按内存里的值换掉的位置记进事件（不含值），事后取质询工件照着再换一遍：质询工件里那一截是「已遮罩」，
和回答时一样。"""
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CALLBACKS)
plays = {"回访": Play([NOTES, PHONES_EAST], "华东的回访备注是 [[v:Q1.r0.备注]]。[[see:Q1]]")}
Models(monkeypatch, name, writer="华东销售额 [[v:Q1.r0.销售额]]。[[see:Q1]]", plays=plays)
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
assert "1390000000" not in r.text
item = r.json()["inquiry"]
assert (item["answer"]["ok"], item["answer"]["error"]) == (False, LEAKED)
notes, phones = item["queries"]
[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, got.text
assert got.json()["content"]["rows"] == [["华东", "客户已遮罩来电改址"], ["华南", "例行回访"]]
assert "1390000000" not in got.text
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{phones['artifact']}")
assert got.json()["content"]["rows"] == [["已遮罩", 1200], ["已遮罩", 300]]
finally:
await engines.invalidate(source_id)
async def test_a_number_written_in_another_table_is_masked_whichever_query_came_first(monkeypatch, tmp_path, client):
"""另一张表（联系方式）的号码列没设遮罩，存的是同一批号码：先查手机号的，查询时就按值换；后查的，事后取质询工件
照回答时记下的位置换。两种先后都整列画成遮罩列。"""
phones = "SELECT buyer_phone AS 手机号, amount AS 金额 FROM sales ORDER BY id"
visits = "SELECT phone_num AS 号码, visits AS 次数 FROM contacts ORDER BY id"
for order in ([phones, visits], [visits, phones]):
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=CONTACTS)
at = order.index(visits) + 1
plays = {"常来": Play(order, f"最常来的客户来过 [[v:Q{at}.r0.次数]] 次。[[see:Q{at}]]")}
Models(monkeypatch, name, writer="华东销售额 [[v:Q1.r0.销售额]]。[[see:Q1]]", plays=plays)
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
assert r.json()["inquiry"]["answer"]["text"] == "最常来的客户来过 5 次。"
query = r.json()["inquiry"]["queries"][at - 1]
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
assert got.status_code == 200, got.text
assert got.json()["content"]["rows"] == [["已遮罩", 5], ["已遮罩", 2]]
assert got.json()["masked"] == ["号码"] and "1390000000" not in got.text
finally:
await engines.invalidate(source_id)
async def test_sql_check_excerpts_in_the_artifact_are_masked_like_the_sql(monkeypatch, masked_shop, client):
"""agent 的查询带着 SQL 检查结果（对照数据目录查出的问题，片段摘自这条 SQL）：质询工件里片段照查询清单的 SQL
换掉遮罩值，给模型的那句改法照抄了片段，不交出去（同证据面板）。"""
@@ -1609,6 +1859,78 @@ async def test_a_customer_looked_up_in_any_shape_never_shows(monkeypatch, tmp_pa
await engines.invalidate(source_id)
#: 同一个号码（13900000001）换了写法：带空格、连字符、点、括号、全角数字或全角标点、+86 / 0086 / 86 前缀
REWRITTEN = [
pytest.param("139-0000-0001", id="连字符"), pytest.param("139 0000 0001", id="空格"),
pytest.param("139.0000.0001", id="点"), pytest.param("(139)0000-0001", id="括号"),
pytest.param("+8613900000001", id="+86"), pytest.param("+86 139 0000 0001", id="+86 加空格"),
pytest.param("0086 139 0000 0001", id="0086"), pytest.param("86-139-0000-0001", id="86"),
pytest.param("１３９０００００００１", id="全角数字"), pytest.param("１３９－００００－０００１", id="全角连字符"),
pytest.param("（139）0000　0001", id="全角括号和空格"),
]
def _shown_anywhere(*texts: str) -> str:
return "\n".join(texts)
@pytest.mark.parametrize("written", REWRITTEN)
async def test_a_masked_number_written_in_another_shape_never_shows(monkeypatch, tmp_path, client, written):
"""验收准备时复现过：模型把号码换个写法写进回答（139-0000-0001、+86 …、全角数字），回答正文、事件、回答文档原样
带着，最后一道关也拦不住，正式运行里永久封存。现在正文里按纯数字写法认出来，换成「已遮罩」。报告取数选了手机号，
号码是从报告自己的查询快照里认出来的。"""
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=SALES)
asked = "SELECT COUNT(*) AS 笔数 FROM sales WHERE buyer_phone = '13900000001'"
plays = {"这位客户": Play([asked], f"客户 {written} 共下了 [[v:Q1.r0.笔数]] 笔。[[see:Q1]]")}
Models(monkeypatch, name, writer="华东首单金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
try:
paused = await pause(name, graph=fetching(name, WITH_PHONE))
[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, "客户 已遮罩 共下了 1 笔。")
[event] = await events(paused.id, "approval.inquiry")
doc = artifact_store.load(item["answer"]["doc_artifact"])
shown = _shown_anywhere(json.dumps(r.json(), ensure_ascii=False), json.dumps(event.data, ensure_ascii=False),
json.dumps(doc, ensure_ascii=False))
assert written not in shown and "1390000000" not in shown
finally:
await engines.invalidate(source_id)
#: 回访记录的备注里写着换了写法的号码（备注没设遮罩）
REWRITTEN_CALLBACKS = SALES + ("CREATE TABLE callbacks (id INTEGER PRIMARY KEY, region TEXT, note TEXT);"
"INSERT INTO callbacks VALUES (1, '华东', '客户 (139)0000-0001 来电改址'), "
"(2, '华南', '例行回访');")
async def test_a_reformatted_number_inside_a_text_cell_is_cut_out_before_the_agent_sees_it(monkeypatch, tmp_path,
client):
"""备注里夹着换了写法的号码（报告取数选了手机号）：开始回答之前就认得出，agent 查备注时那一截已经换掉，引用这一格
回答照常给出；质询工件里同样只换掉那一截。"""
name, source_id = await _source(tmp_path, options={"mask_columns": "buyer_phone"}, script=REWRITTEN_CALLBACKS)
notes = "SELECT region AS 区域, note AS 备注 FROM callbacks ORDER BY id"
plays = {"回访": Play([notes], "华东的回访备注是 [[v:Q1.r0.备注]]。[[see:Q1]]")}
Models(monkeypatch, name, writer="华东首单金额 [[v:Q1.r0.金额]]。[[see:Q1]]", plays=plays)
try:
paused = await pause(name, graph=fetching(name, WITH_PHONE))
[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, "华东的回访备注是 客户 已遮罩 来电改址。")
[query] = item["queries"]
got = await client.get(f"/api/runs/{paused.id}/inquiries/artifacts/{query['artifact']}")
assert got.status_code == 200, got.text
assert got.json()["content"]["rows"] == [["华东", "客户 已遮罩 来电改址"], ["华南", "例行回访"]]
assert "0000-0001" not in got.text and "0000-0001" not in r.text
finally:
await engines.invalidate(source_id)
# --------------------------------------------------------------------------
# 接线：两道审批的竞态、上传表格的版本、设置、开放之后状态又变了、张冠李戴的报告文档
# --------------------------------------------------------------------------
backend/tests/test_inquiry_masking.py
+187
View File
@@ -349,6 +349,119 @@ def test_label_columns_are_traced_to_where_they_are_written(sql, columns, texts)
assert inquiry_mod._label_columns(sql, "sqlite", columns) == {"档位": texts}
#: 派生来源里排在标签前面、和标签同名的列：SQLite 让同名的几列都留着，外层按名字取的是排在前面的那一列（真实等级）
SHADOWED = [
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT *, {BIG_ONE} AS vip_level FROM sales) t", id="* 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT s.*, {BIG_ONE} AS vip_level FROM sales s) t", id="t.* 在前"),
pytest.param(f"WITH t AS (SELECT *, {BIG_ONE} AS vip_level FROM sales) SELECT vip_level AS 档 FROM t",
id="CTE 里 * 在前"),
pytest.param(f"SELECT x.vip_level AS 档 FROM (SELECT * FROM sales s JOIN (SELECT id AS bid, {BIG_ONE} AS vip_level "
"FROM sales) b ON s.id = b.bid) x", id="JOIN 子查询 * 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT *, b.lbl AS vip_level FROM sales s JOIN (SELECT id AS bid, "
f"{BIG_ONE} AS lbl FROM sales) b ON s.id = b.bid) t", id="JOIN 之后 * 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT (vip_level), {BIG_ONE} AS vip_level FROM sales) t",
id="(col) 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT ((vip_level)), {BIG_ONE} AS vip_level FROM sales) t",
id="((col)) 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT vip_level COLLATE NOCASE, {BIG_ONE} AS vip_level FROM sales) t",
id="COLLATE 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT (vip_level) COLLATE BINARY, {BIG_ONE} AS vip_level "
"FROM sales) t", id="(col) COLLATE 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT (vip_level), {BIG_ONE} AS vip_level FROM sales UNION ALL "
f"SELECT vip_level, {BIG_ONE} FROM sales) t", id="多支的第一支 (col) 在前"),
pytest.param(f"SELECT vip_level AS 档 FROM (SELECT vip_level COLLATE NOCASE, {BIG_ONE} AS vip_level FROM sales "
f"UNION ALL SELECT vip_level, {BIG_ONE} FROM sales) t", id="多支的第一支 COLLATE 在前"),
]
@pytest.mark.parametrize("sql", SHADOWED)
def test_a_masked_column_shadowing_a_label_in_a_derived_source_is_not_a_label_column(sql):
"""第 3 轮复验（J2 留下的缺口）：派生来源里 * 带出的遮罩列、没起别名但套了括号或 COLLATE 的遮罩列排在同名的标签
前面，SQLite 交回的是真实等级（1、2、3），以前却按标签认、短码值不收，之后 WHERE vip_level = 1 照原样封存。现在
第一个同名项前面有 * 带出这一列（或说不清带不带）、有去掉括号和 COLLATE 后同名的项，都不算只装着标签的列。"""
plan = plan_of({"source": "shop", "sql": sql, "columns": ["档"], "rows": [[1], [2], [3]]}, {"vip_level"})
assert inquiry_mod._label_columns(sql, "sqlite", ["档"]) == {}
assert {"1", "2", "3"} <= plan.secrets.weak
for level in (1, 2, 3):
assert inquiry_mod._scrub_sql(f"SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = {level}", plan.secrets,
plan.tainted["a1"], "sqlite") == (
"SELECT COUNT(*) AS 人数 FROM sales WHERE vip_level = 已遮罩")
@pytest.mark.parametrize("sql", [
# 标签排在 * 前面：SQLite 取的是排在前面的标签
f"SELECT vip_level AS 档 FROM (SELECT {BIG_ONE} AS vip_level, * FROM sales) t",
# * 只带出另一个派生来源（写了表名）：那里只有标签
f"SELECT x.vip_level AS 档 FROM (SELECT b.* FROM sales s JOIN (SELECT id AS bid, {BIG_ONE} AS vip_level "
"FROM sales) b ON s.id = b.bid) x",
# * 只展开 CTE 里的标签
f"WITH b AS (SELECT {BIG_ONE} AS vip_level FROM sales) SELECT vip_level AS 档 FROM (SELECT * FROM b) t",
])
def test_a_label_ahead_of_a_star_or_alone_behind_one_is_still_a_label_column(sql):
"""收紧之后，标签排在前面、* 只带出派生来源里的标签的，照旧算只装着标签的列。"""
assert inquiry_mod._label_columns(sql, "sqlite", ["档"]) == {"档": {"0", "1"}}
#: 没写表名的 * 连着数据源里的表和装着标签的派生来源
STAR_JOIN = ("SELECT * FROM sales s JOIN (SELECT id AS bid, CASE WHEN amount > 1000 THEN '大单' ELSE '小单' END AS 档 "
"FROM sales) b ON s.id = b.bid")
STAR_COLUMNS = ["id", "region", "amount", "bid", "档"]
def test_a_star_over_a_table_known_to_lack_the_column_keeps_the_label():
"""复核第 1 轮（收紧之后的多遮）：没写表名的 * 展开数据源里的表，以前一律算「说不清带不带这一列」，标签照收成
遮罩值，回答里的「大单」、SQL 里的 '大单' 都换成「已遮罩」。冻结的表结构里查得到这张表、又没有同名的列时，照旧
按派生来源认标签；表里有同名的列、查不到表结构的，照旧不算（宁可多遮）。"""
from app.engine import masking
lacking = {"tables": {"sales": {"columns": [{"name": "id"}, {"name": "region"}, {"name": "amount"}]}}}
holding = {"tables": {"sales": {"columns": [{"name": "id"}, {"name": "档"}, {"name": "amount"}]}}}
assert masking._label_columns(STAR_JOIN, "sqlite", STAR_COLUMNS, masking._schema_tables(lacking)) == {
"档": {"大单", "小单"}}
assert masking._label_columns(STAR_JOIN, "sqlite", STAR_COLUMNS, masking._schema_tables(holding)) == {}
assert masking._label_columns(STAR_JOIN, "sqlite", STAR_COLUMNS) == {}
truncated = {"tables": {"orders": {"columns": [{"name": "id"}]}}, "truncated": True}
assert masking._label_columns(STAR_JOIN, "sqlite", STAR_COLUMNS, masking._schema_tables(truncated)) == {}
# 写了表名的 s.* 展开的是数据源里的表：同样按表结构认它带不带这一列
nested = ("SELECT x.档 FROM (SELECT s.*, b.* FROM sales s JOIN (SELECT id AS bid, CASE WHEN amount > 1000 "
"THEN '大单' ELSE '小单' END AS 档 FROM sales) b ON s.id = b.bid) x")
assert masking._label_columns(nested, "sqlite", ["档"], masking._schema_tables(lacking)) == {"档": {"大单", "小单"}}
assert masking._label_columns(nested, "sqlite", ["档"], masking._schema_tables(holding)) == {}
assert masking._label_columns(nested, "sqlite", ["档"]) == {}
async def test_a_snapshot_reads_labels_with_its_frozen_schema():
"""快照记着查询当时冻结的表结构（schema_artifact）：按它认，* 连着的表没有同名的列时标签不当遮罩值，回答正文、
SQL 里的档名照原样；表结构取不到的照旧不算。"""
from app.core import artifact_store
from app.engine import masking
lacking = {"source": "shop", "tables": {"sales": {"columns": [{"name": "id"}, {"name": "region"},
{"name": "amount"}]}}}
schema = await artifact_store.put_json(lacking, kind="schema_snapshot", run_id="r-labels", node_id="n1")
snapshot = {"source": "shop", "sql": STAR_JOIN, "columns": STAR_COLUMNS, "schema_artifact": schema,
"rows": [[1, "华东", 1200, 1, "大单"], [2, "华南", 300, 2, "小单"]]}
assert masking._snapshot_labels(snapshot, "sqlite") == {"档": {"大单", "小单"}}
plan = plan_of(snapshot, {"amount"})
assert plan.secrets.scrub("这一笔属于大单") == "这一笔属于大单"
sql = "SELECT COUNT(*) FROM (SELECT CASE WHEN amount > 1000 THEN '大单' ELSE '小单' END AS 档 FROM sales) WHERE 档 = '大单'"
assert inquiry_mod._scrub_sql(sql, plan.secrets, plan.tainted["a1"], "sqlite") == sql
assert masking._snapshot_labels({**snapshot, "schema_artifact": "missing"}, "sqlite") == {}
def test_the_tool_boundary_reads_labels_with_the_sources_frozen_schema():
"""工具边界按数据源的表结构认标签：* 连着的表没有同名的列，标签不收成遮罩值，同一份结果里别的文字格里的「大单」照常。"""
from app.engine import masking
payload = {"sql": STAR_JOIN, "columns": [*STAR_COLUMNS, "备注"],
"rows": [[1, "华东", 1200, 1, "大单", "大单客户"], [2, "华南", 300, 2, "小单", "例行回访"]]}
schema = {"tables": {"sales": {"columns": [{"name": "id"}, {"name": "region"}, {"name": "amount"},
{"name": "备注"}]}}}
done = masking.mask_query_result(payload, ["amount"], "sqlite", schema=schema)
assert done.columns == ["amount", "档"] and not {"大单", "小单"} & done.secrets.strong
assert [row[-1] for row in done.payload["rows"]] == ["大单客户", "例行回访"]
# --------------------------------------------------------------------------
# 按值认：门槛、边界
# --------------------------------------------------------------------------
@@ -1144,6 +1257,27 @@ def test_sql_check_excerpts_in_the_artifact_are_masked_like_the_sql():
"sql_excerpt": "SUM(balance)"}]
def test_check_excerpts_follow_the_sql_recorded_at_answer_time():
"""agent 自己的查询快照查询时就遮好了，事后取质询工件时收不到遮罩列里的值：片段、说明照事件里记下的那条 SQL（回答时
按内存里的遮罩值换过）换——那条里整个换成「已遮罩」的字面量、换掉的那一截（LIKE 的 %…%），片段、说明里一样换。"""
sql = ("SELECT COUNT(*) AS 笔数 FROM sales WHERE buyer_phone = '13900000001' AND age = 25 "
"AND note LIKE '%13900000002%'")
recorded = "SELECT COUNT(*) AS 笔数 FROM sales WHERE buyer_phone = '已遮罩' AND age = 已遮罩 AND note LIKE '%已遮罩%'"
assert inquiry_mod._recorded_swaps(sql, recorded, "sqlite") == ({"13900000001", "25"}, {"13900000002"})
assert inquiry_mod._recorded_swaps(sql, sql, "sqlite") == (set(), set())
assert inquiry_mod._recorded_swaps(sql, recorded + " LIMIT 5", "sqlite") == (set(), set()), "对不上就不认"
checks = [{"code": "c1", "level": "warning", "table": "sales", "message": "没有码值「25」",
"sql_excerpt": "buyer_phone = '13900000001' AND age = 25"},
{"code": "c2", "level": "info", "table": "sales", "message": "按备注筛选",
"sql_excerpt": "note LIKE '%13900000002%'"}]
whole, pieces = inquiry_mod._recorded_swaps(sql, recorded, "sqlite")
assert inquiry_mod._masked_checks(checks, sql, inquiry_mod._Secrets(), set(), "sqlite", swapped=whole,
pieces=pieces) == [
{"code": "c1", "level": "warning", "table": "sales", "message": "没有码值「已遮罩」",
"sql_excerpt": "buyer_phone = '已遮罩' AND age = 已遮罩"},
{"code": "c2", "level": "info", "table": "sales", "message": "按备注筛选", "sql_excerpt": "note LIKE '%已遮罩%'"}]
@pytest.mark.parametrize("sql, excerpt, shown", [
# 片段只是一段表达式，单独拿出来认不出在找哪一个值：按整条 SQL 里换掉的字面量换
("SELECT SUM(CASE age WHEN 25 THEN balance END) AS 余额 FROM accounts",
@@ -1299,3 +1433,56 @@ def test_hundreds_of_band_edges_are_read_within_a_budget():
assert "已遮罩" in inquiry_mod._scrub_sql(bins(600), every, {"age"}, "sqlite")
elapsed = time.perf_counter() - started
assert elapsed < 5, f"用了 {elapsed:.2f} 秒"
# --------------------------------------------------------------------------
# 回答时才认得出的值：记下位置，事后取质询工件时照着再换一遍
# --------------------------------------------------------------------------
def test_values_known_only_at_answer_time_are_recorded_as_positions_and_replayed():
"""agent 先查了备注（那时还不知道备注里的号码是遮罩值，查询时原样交回、原样存下），后来又查出了手机号（查询时
遮掉，值只在内存里）。回答时按内存里的值换掉的地方记成位置（不含值）；事后取质询工件时内存里的值已经不在，照着
位置再换一遍，和回答时看到的一样。"""
from app.engine import masking
note = "客户13900000001来电改址，备用 139-0000-0001"
snapshot = {"columns": ["区域", "备注", "号码", "13900000001 的单"],
"rows": [["华东", note, "13900000001", 1], ["华南", "例行回访", "13800000000", 2]]}
seen = secrets_of({"columns": ["phone"], "rows": [["13900000001"]]}, {"phone"})
edits = masking.view_edits(snapshot, set(), seen, inquiry_mod._Secrets())
first, second = note.index("13900000001"), note.index("139-0000-0001")
assert edits == {"cells": [[0, 2]], "spans": [[0, 1, first, first + 11], [0, 1, second, second + 13]],
"names": [[3, 0, 11]], "masked": ["号码"]}
assert "1390000000" not in json.dumps(edits, ensure_ascii=False), "只记位置，不记值"
replayed = masking.replay_edits(snapshot, [edits])
assert replayed["columns"] == ["区域", "备注", "号码", f"{MASKED} 的单"]
assert replayed["rows"] == [["华东", f"客户{MASKED}来电改址，备用 {MASKED}", MASKED, 1],
["华南", "例行回访", "13800000000", 2]]
# 不靠内存里的值也认得出的（报告快照里就有），事后照样认得出，不记
assert masking.view_edits(snapshot, set(), seen, seen) is None
# 记下的位置对不上（越界、坏形状）的跳过，不出错
broken = {"cells": [[9, 9], "x"], "spans": [[0, 1, 50, 60], [0, 9, 0, 1]], "names": [[7, 0, 1]], "masked": [3]}
assert masking.replay_edits(snapshot, [broken, None, "x"]) == snapshot
def test_the_plan_records_what_only_the_values_seen_by_the_agent_masked():
"""_settle_plan 带上 agent 查询时遮掉的值（seen）时，另记下只靠它们才换掉的位置（replay）；不带 seen 的（事后取
质询工件）不记。"""
notes = {"source": "shop", "sql": "SELECT region AS 区域, note AS 备注 FROM callbacks", "columns": ["区域", "备注"],
"rows": [["华东", "客户13900000001来电改址"], ["华南", "例行回访"]]}
seen = secrets_of({"columns": ["phone"], "rows": [["13900000001"]]}, {"phone"})
plan = inquiry_mod._settle_plan({"a1": notes}, {"a1": ({"buyer_phone"}, [])}, [], {"shop": "sqlite"}, [], [], {},
frozenset(), seen)
assert plan.replay == {"a1": {"cells": [], "spans": [[0, 1, 2, 13]], "names": [], "masked": []}}
later = inquiry_mod._settle_plan({"a1": notes}, {"a1": ({"buyer_phone"}, [])}, [], {"shop": "sqlite"}, [], [], {},
frozenset({"a1"}))
assert later.replay == {} and later.views["a1"]["rows"][0][1] == "客户13900000001来电改址"
def test_values_the_agent_saw_survive_a_report_that_fills_the_cap():
"""报告快照里的遮罩值收满了上限（SECRETS_MAX）：agent 查询时遮掉的值先收，照样认得出（以前先收报告的，满了就丢）。"""
report = {"columns": ["buyer_phone"], "rows": [[str(13800000000 + i)] for i in range(inquiry_mod.SECRETS_MAX)]}
seen = secrets_of({"columns": ["phone"], "rows": [["19914742785"]]}, {"phone"})
plan = inquiry_mod._settle_plan({}, {}, [(report, {"buyer_phone"})], {}, [], [], {}, frozenset(), seen)
assert plan.secrets.hit("19914742785")
backend/tests/test_masking_reformatted_numbers.py
+209
View File
@@ -0,0 +1,209 @@
"""换了写法的号码（engine/masking.py 的 _Secrets）：按值认得出的长数字（SECRET_DIGITS 位以上）被写成带空格、连字符、
点、括号、全角数字或全角标点、带 +86 / 0086 / 86 前缀的样子，模型写的正文（审批质询的回答）里照样认得出，换成
「已遮罩」；最后一道关（inquiry._leaked）、质询工件里文字格夹着的号码同样认得出。
验收准备时复现过：模型把号码写成 199-1474-2785、199 1474 2785、199.1474.2785、(199)1474-2785、+8619914742785、
全角数字，回答正文、事件、回答文档原样带着，最后一道关也拦不住，正式运行里永久封存。
认法是在正文里找允许夹分隔符的数字串，规范化（全角转半角、去分隔符、去国家码）后在「遮罩值的纯数字写法」集合里
查——集合查找，不现编正则。日期、小数、更长的数里的一截、连字符连着的编号照旧不动。
数据都是虚构的：手机号、金额。
"""
from __future__ import annotations
import time
import pytest
from app.engine import inquiry as inquiry_mod
from app.engine.masking import _masked_view, _Secrets
MASKED = "已遮罩"
PHONE = "19914742785"
#: 同一个号码的各种写法（换了写法的号码）
REFORMATTED = [
pytest.param("199-1474-2785", id="连字符"),
pytest.param("199 1474 2785", id="空格"),
pytest.param("199.1474.2785", id="点"),
pytest.param("(199)1474-2785", id="括号"),
pytest.param("(199) 1474 2785", id="括号加空格"),
pytest.param("+8619914742785", id="+86"),
pytest.param("+86 199 1474 2785", id="+86 加空格"),
pytest.param("(+86)19914742785", id="(+86)"),
pytest.param("0086 199 1474 2785", id="0086"),
pytest.param("0086-19914742785", id="0086 连字符"),
pytest.param("86-199-1474-2785", id="86"),
pytest.param("8619914742785", id="86 连写"),
pytest.param("１９９１４７４２７８５", id="全角数字"),
pytest.param("１９９－１４７４－２７８５", id="全角数字和连字符"),
pytest.param("199　1474　2785", id="全角空格"),
pytest.param("199．1474．2785", id="全角点"),
pytest.param("（199）1474-2785", id="全角括号"),
pytest.param("＋86 199 1474 2785", id="全角加号"),
]
def secrets_of(values: list, *, hidden: str = "phone") -> _Secrets:
secrets = _Secrets()
secrets.add({"columns": [hidden], "rows": [[v] for v in values]}, {hidden})
return secrets
@pytest.mark.parametrize("written", REFORMATTED)
@pytest.mark.parametrize("stored", [PHONE, int(PHONE), "199-1474-2785"], ids=["文本", "整数", "存的就带连字符"])
def test_a_long_number_in_another_shape_is_masked_in_prose(stored, written):
secrets = secrets_of([stored])
assert secrets.scrub(f"金额最高的一笔来自 {written}，共 3 笔", prose=True) == f"金额最高的一笔来自 {MASKED}，共 3 笔"
assert secrets.scrub(f"客户{written}来电", prose=True) == f"客户{MASKED}来电"
@pytest.mark.parametrize("text, shown", [
("来自（19914742785），共 3 笔", f"来自（{MASKED}），共 3 笔"),
("来自(+86 199 1474 2785)，共 3 笔", f"来自({MASKED})，共 3 笔"),
("电话（+86）199 1474 2785。", f"电话{MASKED}。"),
("（电话 199-1474-2785）", f"（电话 {MASKED}）"),
])
def test_brackets_around_a_number_are_kept_and_brackets_inside_it_are_replaced(text, shown):
"""整个号码外面套着的一对括号是标点，留着；号码里的括号（区号、国家码）和号码一起换掉。"""
assert secrets_of([PHONE]).scrub(text, prose=True) == shown
@pytest.mark.parametrize("text", [
# 日期、时间、小数、千分位、更长的数里的一截、连字符连着的编号、别的号码
"2026-10-06 19:14 下单", "金额 19914.74 元", "共 1,991,474 笔", "号码 1991474278512", "批次 19914742785-2",
"199-1474-2786", "单号 2026199-1474-2785", "共 199 笔 1474 元 2785 次",
])
def test_other_numbers_and_number_like_text_are_left_alone(text):
secrets = secrets_of([PHONE])
assert secrets.scrub(text, prose=True) == text
def test_reformatted_numbers_are_only_read_in_prose_and_text_cells():
"""SQL 里 199-1474-2785 是一个算式，不当成号码换（SQL 里的号码按字面量的用法换，见 inquiry._scrub_sql）。"""
secrets = secrets_of([PHONE])
assert secrets.scrub("SELECT 199-1474-2785 AS x") == "SELECT 199-1474-2785 AS x"
assert secrets.scrub("备注 199-1474-2785", loose=True) == f"备注 {MASKED}"
def test_a_text_cell_holding_a_reformatted_number_is_masked_whole_or_in_part():
"""快照里别的列（没设遮罩）写着换了写法的号码：整格就是号码的整格遮；质询工件里夹在一段文字里的只换那一截。"""
secrets = secrets_of([PHONE])
snapshot = {"columns": ["联系方式", "备注"], "rows": [["199-1474-2785", "客户 (199)1474-2785 来电改址"],
["138-0000-0000", "例行回访"]]}
view, masked = _masked_view(snapshot, set(), secrets)
assert view["rows"] == [[MASKED, "客户 (199)1474-2785 来电改址"], ["138-0000-0000", "例行回访"]]
assert masked == ["联系方式"]
view, masked = _masked_view(snapshot, set(), secrets, embedded=True)
assert view["rows"] == [[MASKED, f"客户 {MASKED} 来电改址"], ["138-0000-0000", "例行回访"]]
assert masked == ["联系方式"]
@pytest.mark.parametrize("written", REFORMATTED)
def test_the_last_gate_stops_an_answer_still_showing_a_reformatted_number(written):
"""组装之后还认得出（渲染进单元格引用的一格文字里夹着换了写法的号码）：不交出去。正文里的在组装之前就换掉了。"""
secrets = secrets_of([PHONE])
entry = {"alias": "Q1", "kind": "query", "artifact": "a1", "source": "shop", "columns": ["备注"], "rows": 1,
"truncated": False}
views = {"a1": _masked_view({"columns": ["备注"], "rows": [[f"客户{written}来电改址"]]}, set(), secrets)[0]}
with pytest.raises(inquiry_mod._Leak):
inquiry_mod._compose("回访备注是 [[v:Q1.r0.备注]]。", {"Q1": entry}, run_id="r1", views=views, secrets=secrets,
allow=[])
doc = inquiry_mod._compose(f"来电的是 {written}。", {}, run_id="r1", views={}, secrets=secrets, allow=[])
assert doc["markdown"] == f"来电的是 {MASKED}。"
assert inquiry_mod._leaked({"markdown": f"来电的是 {written}。", "blocks": []}, secrets)
def test_a_number_the_asker_wrote_in_another_shape_is_left_alone():
"""提问人自己写的号码（换了写法也一样）已经摆在读者眼前：回答里照原样，别的号码照遮。"""
secrets = secrets_of([PHONE, "13800000000"])
secrets.protect(["199-1474-2785 这位客户还买过别的吗？"])
assert secrets.scrub("199 1474 2785 另有 138-0000-0000 的单", prose=True) == f"199 1474 2785 另有 {MASKED} 的单"
def test_a_threshold_left_in_the_sql_is_shown_in_the_prose_whatever_the_shape():
"""SQL 里照原样留着的阈值，回答正文里复述不遮，换了写法也一样（只免正文）。"""
secrets = secrets_of([20000, 35000], hidden="amount")
secrets.show({"20000"})
assert secrets.scrub("金额不低于 20 000 元，最高一笔 35 000 元", prose=True) == f"金额不低于 20 000 元，最高一笔 {MASKED} 元"
def test_many_masked_numbers_are_matched_without_stalling():
"""遮罩值上万个：集合查找，不现编上万个分支的正则。"""
secrets = secrets_of([str(19900000000 + i) for i in range(20_000)])
text = "；".join(f"客户 199-0000-{i:04d} 来电" for i in range(0, 4000, 7)) * 3
began = time.perf_counter()
out = secrets.scrub(text, prose=True)
assert time.perf_counter() - began < 2.0
assert "199-0000" not in out and out.count(MASKED) == text.count("199-0000")
#: 复核第 1 轮：别的空白、减号作分隔符的写法（模型输出里不换行空格、窄空格并不少见）
OTHER_SEPARATORS = [
pytest.param("199 1474 2785", id="不换行空格"),
pytest.param("199 1474 2785", id="窄不换行空格"),
pytest.param("199 1474 2785", id="细空格"),
pytest.param("199 1474 2785", id="发丝空格"),
pytest.param("199\t1474\t2785", id="制表符"),
pytest.param("199−1474−2785", id="减号"),
pytest.param("199 1474 2785", id="四个空格"),
pytest.param("199 - 1474 - 2785", id="连字符两边带空格"),
pytest.param("+86 199 1474 2785", id="+86 加不换行空格"),
]
@pytest.mark.parametrize("written", OTHER_SEPARATORS)
def test_numbers_split_by_other_spaces_and_minus_signs_are_masked_too(written):
secrets = secrets_of([PHONE])
assert secrets.scrub(f"客户 {written} 来电", prose=True) == f"客户 {MASKED} 来电"
assert inquiry_mod._leaked({"markdown": f"来电的是 {written}。", "blocks": []}, secrets)
assert secrets.hit(written), "整格就是这个号码的，格子按值认"
def test_a_line_break_still_ends_a_number():
"""不跨行：两行各写一截的不拼成一个号码。"""
secrets = secrets_of([PHONE])
assert secrets.scrub("编号 199\n1474 2785", prose=True) == "编号 199\n1474 2785"
@pytest.mark.parametrize("sql, sealed", [
pytest.param("SELECT COUNT(*) AS 笔数 FROM callbacks WHERE note LIKE '%199-1474-2785%'",
f"SELECT COUNT(*) AS 笔数 FROM callbacks WHERE note LIKE '%{MASKED}%'", id="LIKE 里的一截"),
pytest.param("SELECT COUNT(*) AS 笔数 FROM callbacks WHERE note = '客户 199 1474 2785 来电'",
f"SELECT COUNT(*) AS 笔数 FROM callbacks WHERE note = '客户 {MASKED} 来电'", id="字符串里夹着"),
pytest.param('SELECT note AS "199-1474-2785 的备注" FROM callbacks',
f'SELECT note AS "{MASKED} 的备注" FROM callbacks', id="带引号的列名"),
pytest.param("SELECT COUNT(*) AS 笔数 FROM callbacks /* 客户 199-1474-2785 */",
f"SELECT COUNT(*) AS 笔数 FROM callbacks /* 客户 {MASKED} */", id="块注释"),
pytest.param("SELECT COUNT(*) AS 笔数 FROM callbacks -- 客户 199 1474 2785\n WHERE id = 1",
f"SELECT COUNT(*) AS 笔数 FROM callbacks -- 客户 {MASKED}\n WHERE id = 1", id="行注释"),
])
def test_a_reformatted_number_in_sql_strings_names_and_comments_is_masked(sql, sealed):
"""SQL 里字符串、带引号的名字、注释里换了写法的号码也认（以前只有整个字面量就是号码的才换）。"""
assert inquiry_mod._scrub_sql(sql, secrets_of([PHONE]), set(), "sqlite") == sealed
def test_a_bare_number_expression_in_sql_is_still_left_alone():
"""不带引号的 199-1474-2785 在 SQL 里是一个算式（三个数相减），照旧不按号码认。"""
sql = "SELECT 199-1474-2785 AS x FROM callbacks"
assert inquiry_mod._scrub_sql(sql, secrets_of([PHONE]), set(), "sqlite") == sql
def test_a_threshold_matching_a_number_swapped_in_a_string_is_not_kept():
"""字符串里按换了写法认出、换掉的号码，同一条 SQL 里写成阈值的同一个数也不留（不然阈值把它露了）。"""
sql = f"SELECT COUNT(*) AS 笔数 FROM callbacks WHERE note LIKE '%199-1474-2785%' AND phone_no >= {PHONE}"
sealed = inquiry_mod._scrub_sql(sql, secrets_of([PHONE]), set(), "sqlite")
assert PHONE not in sealed and "1474" not in sealed
def test_a_reformatted_number_in_a_column_name_is_masked_like_the_sql():
"""agent 把换了写法的号码写进列名（AS "139-0000-0001 的单"）：质询工件的列名、回答里引用这一列的标记照 SQL 里带引号
的名字一样换，三处对得上。"""
secrets = secrets_of(["13900000001"])
snapshot = {"columns": ["139-0000-0001 的单"], "rows": [[3]]}
view, _ = _masked_view(snapshot, set(), secrets)
assert view["columns"] == [f"{MASKED} 的单"]
assert inquiry_mod._scrub_sql('SELECT COUNT(*) AS "139-0000-0001 的单" FROM sales', secrets, set(), "sqlite") == (
f'SELECT COUNT(*) AS "{MASKED} 的单" FROM sales')
assert inquiry_mod._scrubbed_source("共 [[v:Q1.r0.139-0000-0001 的单]] 单。", secrets) == f"共 [[v:Q1.r0.{MASKED} 的单]] 单。"
backend/tests/test_masking_star_renames.py
+378
View File
@@ -0,0 +1,378 @@
"""按 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", "合计"}
# --------------------------------------------------------------------------
# 写了库名的表、过期的表结构（复核第 3 轮）
# --------------------------------------------------------------------------
#: 同名的两张表：主库的 sales 和另一个库（archive）里的 sales，号码、区域换了次序。表结构探查的是主库
TWO_SALES = ("CREATE TABLE sales (id INTEGER, region TEXT, buyer_phone TEXT, buyer_name TEXT, amount INTEGER, "
"note TEXT);"
"INSERT INTO sales VALUES (1, '华东', '13900000001', '张伟', 1200, '改址');"
"CREATE TABLE archive.sales (id INTEGER, buyer_phone TEXT, region TEXT, buyer_name TEXT, amount INTEGER, "
"note TEXT);"
"INSERT INTO archive.sales VALUES (7, '13900000007', '华北', '王芳', 500, '回访');")
SALES_LAYOUT = ["id", "region", "buyer_phone", "buyer_name", "amount", "note"]
@pytest.fixture
def two_sales():
conn = sqlite3.connect(":memory:")
conn.execute("ATTACH DATABASE ':memory:' AS archive")
conn.executescript(TWO_SALES)
yield conn
conn.close()
def probed(schema: str | None = None) -> dict:
"""探查主库得到的表结构（schema_cache 的样子）：没设 schema 的探查当前库，qualified 就是表名。"""
meta = {"qualified": f"{schema}.sales" if schema else "sales", "schema": schema,
"columns": [{"name": c} for c in SALES_LAYOUT]}
return {"tables": {"sales": meta}, "schema": schema}
def run_two(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, MASKED, "sqlite", schema=schema)
@pytest.mark.parametrize("sql", [
pytest.param("WITH t(a, b, c, d, e, f) AS (SELECT * FROM archive.sales) SELECT b AS j, a AS n FROM t",
id="列名表套 * 展开另一个库里的同名表"),
pytest.param("SELECT 0 AS a, '' AS b, '' AS c, '' AS d, 0 AS e, '' AS f WHERE 0 UNION ALL "
"SELECT * FROM archive.sales", id="UNION 后一支的 * 展开另一个库里的同名表"),
pytest.param("WITH t AS (SELECT * FROM sales UNION ALL SELECT * FROM archive.sales) SELECT region AS 区域 FROM t",
id="派生来源里 UNION 后一支展开另一个库里的同名表"),
])
def test_a_table_in_another_schema_is_not_counted_by_the_probed_one(two_sales, sql):
"""表结构探查的是主库：查询写了别的库名的同名表（archive.sales，号码排在第二列），不能按主库那张表的列序数位置，
按数不清处理，* 展开落到的那一段整段遮。以前按表名查表结构，号码落进的那一列被当成区域，原样交给模型、落库。"""
done = run_two(two_sales, sql, probed())
assert not _leaks(done), done.payload["rows"]
def test_a_table_named_with_its_probed_schema_is_still_counted(two_sales):
"""库名和表结构记下的对得上（探查时设了 schema、qualified 带着它）：照样按表结构数位置，只遮号码那一列。没写库名的
照旧按表名查。库名对不上表结构记下的（这里探查时没设 schema，查询写了 main.）按数不清处理，整段遮（多遮）。"""
precise = "WITH t(a, b, c, d, e, f) AS (SELECT * FROM {}) SELECT c AS j, a AS n FROM t"
assert run_two(two_sales, precise.format("main.sales"), probed("main")).columns == ["j"]
assert run_two(two_sales, precise.format("sales"), probed()).columns == ["j"]
assert set(run_two(two_sales, precise.format("main.sales"), probed()).columns) == {"j", "n"}
def test_schema_columns_are_also_keyed_by_the_probed_schema():
"""表结构按「库名.表名」另记一份（库名取 schema、qualified 的前缀）；写了别的库名的查不到。"""
tables = masking._schema_columns(probed("public"))
assert tables["sales"] == tables["public.sales"] == SALES_LAYOUT and "archive.sales" not in tables
assert "public.sales" not in masking._schema_columns(probed())
assert "archive.sales" not in (masking._schema_tables(probed("public")) or {})
#: 一月的表结构是早先探查的（那时还没有 extra 这一列），表里现在多了一列，排在第二；二月那张号码排在第二列
STALE = ("CREATE TABLE jan (id INTEGER, extra TEXT, region TEXT, buyer_phone TEXT, amount INTEGER);"
"CREATE TABLE feb (id INTEGER, buyer_phone TEXT, region TEXT, amount INTEGER, extra TEXT);"
"INSERT INTO jan VALUES (1, '加急', '华东', '13900000001', 1200);"
"INSERT INTO feb VALUES (2, '13900000002', '华南', 300, '普通');")
@pytest.fixture
def stale():
conn = sqlite3.connect(":memory:")
conn.executescript(STALE)
yield conn
conn.close()
def layout(**tables: list[str]) -> dict:
return {"tables": {name: {"qualified": name, "columns": [{"name": c} for c in columns]}
for name, columns in tables.items()}}
#: 过期的表结构：一月那张少了 extra（列数和表里实际的不一样）
OLD_JAN = layout(jan=["id", "region", "buyer_phone", "amount"], feb=["id", "buyer_phone", "region", "amount", "extra"])
#: 准的表结构
NOW = layout(jan=["id", "extra", "region", "buyer_phone", "amount"],
feb=["id", "buyer_phone", "region", "amount", "extra"])
@pytest.mark.parametrize("sql", [
pytest.param("WITH t AS (SELECT * FROM jan UNION ALL SELECT id, buyer_phone, region, amount, extra FROM feb) "
"SELECT extra AS 备注 FROM t", id="后一支明写各列"),
pytest.param("WITH t AS (SELECT * FROM jan UNION ALL SELECT * FROM feb) SELECT extra AS 备注 FROM t",
id="后一支也是 *"),
])
def test_a_first_branch_counted_by_a_stale_schema_is_not_trusted(stale, sql):
"""UNION 第一支是 *、按表结构数出的列名（表结构过期了，少一列）和后一支的列数对不上：不按它对列名，* 所在的位置按认不出
算，整份遮。以前照过期的列名表对，二月的号码落进的「extra」没认出来，原样交回。"""
cursor = stale.execute(sql)
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
for schema in (OLD_JAN, None, NOW):
done = masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=schema)
assert done.columns == ["备注"] and not _leaks(done), (schema, done.payload["rows"])
def test_a_first_branch_counted_by_an_accurate_schema_still_places_the_values(stale):
"""表结构准、每一支数得出的列数都对得上：照样按第一支的列名对上位置，只遮号码落进的那一列，金额照常。"""
sql = ("WITH t AS (SELECT * FROM jan UNION ALL SELECT id, buyer_phone, region, buyer_phone, amount FROM feb) "
"SELECT region AS 区域, amount AS 金额 FROM t")
cursor = stale.execute(sql)
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
assert masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=NOW).columns == []
moved = sql.replace("SELECT id, buyer_phone, region,", "SELECT id, region, buyer_phone,")
cursor = stale.execute(moved)
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()],
"sql": moved}
done = masking.mask_query_result(payload, ["buyer_phone"], "sqlite", schema=NOW)
assert done.columns == ["区域"] and [row[1] for row in done.payload["rows"]] == [1200, 300]
backend/tests/test_masking_table_functions.py
+161
View File
@@ -0,0 +1,161 @@
"""按 SQL 认取自遮罩列的结果列（engine/masking.py 的 _derived_hidden）：FROM、JOIN、LATERAL 里的表值函数，汇总里的 CASE。
核查 agent 的查询在工具边界按这套认法整列遮掉（mask_query_result），审批质询、双盲复算也用它：
- 表值函数（SQLite 的 json_each / json_tree，Postgres 的 unnest、jsonb_each_text、regexp_split_to_table、LATERAL
VALUES，MySQL 的 JSON_TABLE）把遮罩列拆开时，它出的列都算取自遮罩列：遮罩列是 JSON 时，用 json_each 拆开是最自然的
写法，以前只追 SELECT 项、子查询和 *，这些写法的原值原样交给模型、落进快照。出列名认得全的（常见的几个函数、写了
列名表的）按名字认，认不全的整份遮
- 汇总里的 CASE / IIF 只在条件里用到遮罩列的（SUM(CASE WHEN buyer_phone IS NOT NULL THEN amount ELSE 0 END)），同
写在 WHERE 里，不算；结果的某一支取的是遮罩列、不汇总的逐行 CASE 照旧算
数据都是虚构的：手机号、区域、金额。
"""
from __future__ import annotations
import pytest
from app.engine import masking
#: 设了遮罩的列：手机号，和装着手机号、姓名的 JSON 列
MASKED = {"buyer_phone", "profile"}
SALES = ["id", "region", "buyer_phone", "amount", "profile"]
#: SQLite 的 json_each / json_tree 出的列
JSON_EACH = ["key", "value", "type", "atom", "id", "parent", "fullkey", "path"]
def hidden_of(sql: str, columns: list[str], kind: str = "sqlite", masked: set[str] = MASKED) -> set[str]:
snapshot = {"source": "shop", "sql": sql, "columns": columns, "rows": []}
return masking._derived_hidden(snapshot, set(masked), kind)[0]
@pytest.mark.parametrize("kind, sql, columns, shown", [
pytest.param("sqlite", "SELECT j.key, j.value FROM sales, json_each(sales.profile) AS j", ["key", "value"],
["key", "value"], id="json_each 拆遮罩的 JSON 列"),
pytest.param("sqlite", "SELECT t.atom AS 值 FROM sales, json_tree(sales.profile) AS t WHERE t.type = 'text'",
["值"], ["值"], id="json_tree"),
pytest.param("sqlite", "SELECT value FROM sales, json_each(json_array(buyer_phone))", ["value"], ["value"],
id="json_each 套 json_array"),
pytest.param("sqlite", "SELECT r.value AS 号码 FROM sales s JOIN json_each(s.profile) r ON 1 = 1", ["号码"],
["号码"], id="JOIN 里的 json_each"),
pytest.param("sqlite", "SELECT value FROM json_each((SELECT json_group_array(buyer_phone) FROM sales))",
["value"], ["value"], id="参数是子查询"),
pytest.param("sqlite", "SELECT c.value AS 号码 FROM (SELECT json_group_array(buyer_phone) AS arr FROM sales) t, "
"json_each(t.arr) c", ["号码"], ["号码"], id="参数取自派生来源里改了名字的遮罩列"),
pytest.param("sqlite", "SELECT * FROM sales, json_each(sales.profile)", [*SALES, *JSON_EACH],
["buyer_phone", "profile", "key", "value", "atom", "fullkey", "path"], id="* 带出 json_each 的列"),
pytest.param("sqlite", "SELECT x FROM (SELECT j.value AS x FROM sales, json_each(sales.profile) j) t", ["x"],
["x"], id="子查询里拆开再改名"),
pytest.param("postgres", "SELECT u FROM sales, unnest(ARRAY[buyer_phone]) AS u", ["u"], ["u"], id="unnest"),
pytest.param("postgres", "SELECT x.p FROM sales s CROSS JOIN LATERAL unnest(ARRAY[s.buyer_phone]) AS x(p)",
["p"], ["p"], id="LATERAL unnest 带列名表"),
pytest.param("postgres", "SELECT x FROM sales, LATERAL unnest(ARRAY[buyer_phone]) AS x", ["x"], ["x"],
id="LATERAL unnest"),
pytest.param("postgres", "SELECT u FROM sales CROSS JOIN LATERAL (VALUES (buyer_phone)) AS v(u)", ["u"], ["u"],
id="LATERAL VALUES"),
pytest.param("postgres", "SELECT d FROM sales, regexp_split_to_table(buyer_phone, '-') AS d", ["d"], ["d"],
id="regexp_split_to_table"),
pytest.param("postgres", "SELECT e.key, e.value FROM sales, jsonb_each_text(profile) AS e", ["key", "value"],
["key", "value"], id="jsonb_each_text"),
pytest.param("postgres", "SELECT x FROM sales, unnest(string_to_array(buyer_phone, ',')) AS x", ["x"], ["x"],
id="unnest 套 string_to_array"),
pytest.param("postgres", "SELECT * FROM sales, unnest(ARRAY[buyer_phone])", [*SALES, "unnest"],
["buyer_phone", "profile", "unnest"], id="没起别名的 unnest 按函数名出列"),
pytest.param("mysql", "SELECT j.p FROM sales s, JSON_TABLE(JSON_ARRAY(s.buyer_phone), '$[*]' "
"COLUMNS (p VARCHAR(20) PATH '$')) j", ["p"], ["p"], id="JSON_TABLE"),
])
def test_a_table_function_over_a_masked_column_hides_what_it_returns(kind, sql, columns, shown):
assert {c.lower() for c in shown} <= hidden_of(sql, columns, kind)
def test_columns_a_table_function_does_not_return_stay_visible():
"""拆的是遮罩列，别的列（区域、金额）照常；拆的不是遮罩列的，一列都不遮。"""
hidden = hidden_of("SELECT s.region, s.amount, j.value AS 值 FROM sales s, json_each(s.profile) j",
["region", "amount", "值"])
assert "值" in hidden and not {"region", "amount"} & hidden
assert hidden_of("SELECT s.region, j.value AS 标签 FROM sales s, json_each(s.tags) j", ["region", "标签"]) == MASKED
def test_a_table_function_that_cannot_be_read_hides_the_whole_snapshot():
"""不认得的函数拆遮罩列、又没写列名表：说不清出哪几列，整份快照的列都遮（同 SQL 解析不了）。"""
columns = ["region", "piece"]
assert hidden_of("SELECT s.region, t.piece FROM sales s, my_split(s.buyer_phone) t", columns, "postgres") == {
*MASKED, "region", "piece"}
@pytest.mark.parametrize("sql", [
"SELECT SUM(CASE WHEN buyer_phone IS NOT NULL THEN amount ELSE 0 END) AS total FROM sales",
"SELECT SUM(CASE buyer_phone WHEN '13900000001' THEN amount END) AS total FROM sales",
"SELECT AVG(IIF(buyer_phone LIKE '139%', amount, NULL)) AS total FROM sales",
"SELECT MAX(CASE WHEN buyer_phone IS NULL THEN 0 ELSE amount END) AS total FROM sales",
"SELECT SUM(CASE WHEN buyer_phone IS NULL THEN 0 ELSE CASE WHEN region = '华东' THEN amount END END) AS total "
"FROM sales",
])
def test_a_masked_column_only_in_the_condition_of_an_aggregated_case_is_like_a_filter(sql):
"""值全部来自金额，手机号只决定这一行算不算进来：同写在 WHERE 里，不算取自遮罩列。"""
assert hidden_of(sql, ["total"]) == MASKED
assert hidden_of(sql.replace("CASE WHEN buyer_phone IS NOT NULL THEN amount ELSE 0 END", "amount")
+ " WHERE buyer_phone IS NOT NULL", ["total"]) == MASKED
@pytest.mark.parametrize("sql", [
# 结果的某一支取的就是遮罩列
"SELECT SUM(CASE WHEN region = '华东' THEN buyer_phone END) AS total FROM sales",
"SELECT MAX(IIF(amount > 1000, buyer_phone, NULL)) AS total FROM sales",
# 不汇总的逐行 CASE：每一行各自落在哪一支
"SELECT CASE WHEN buyer_phone IS NULL THEN 0 ELSE amount END AS total FROM sales",
])
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"] == [["已遮罩"], ["已遮罩"]]
backend/tests/test_masking_whole_rows.py
+244
View File
@@ -0,0 +1,244 @@
"""按 SQL 认取自遮罩列的结果列（engine/masking.py 的 _derived_hidden）：整行、函数参数里的 *、列名表把遮罩列带进别的列名
（复核第 3 轮）。
Postgres 里不带表名、名字是 FROM 里一个来源的列引用指它的整行（row_to_json(s)、json_agg(t)、SELECT s、s::text），函数
参数里的 s.*（to_json(s.*)、concat_ws(',', s.*)、ROW(s.*)）也把整行交给函数；数据源里的表可以带列名表（FROM sales AS
s(i, r, p)，只改前几列的名字），TABLE sales 等于 SELECT * FROM sales；Oracle 的 JSON_OBJECT(*) 把整行拼成 JSON，UNPIVOT、
PIVOT、MATCH_RECOGNIZE 把遮罩列的值转进它们自己起名的列。这些写法以前都认不出：遮罩列的值落进了别的列名，原样交给核查
agent、落进快照。现在整行、函数参数里的 * 按 * 展开的来源认（数据源里的表一律算带出了遮罩列），列名表按位置对上（表结构
查得到这张表的只遮号码改了名字的那一列），TABLE x 当 SELECT * FROM x，PIVOT 这类读到遮罩列的整份遮。
这里没有真的 Postgres、Oracle：结果按它们交回的样子造（行转成 JSON、记录转成文字），只看遮完的结果里还有没有原值。数据
都是虚构的：区域、手机号、姓名、金额。
"""
from __future__ import annotations
import json
import sqlite3
import pytest
from app.engine import masking
MASKED_WORD = "已遮罩"
MASKED = ["buyer_phone", "buyer_name"]
SECRETS = ("1390000000", "张伟", "李娜")
LAYOUT = ["id", "region", "buyer_phone", "buyer_name", "amount", "note"]
#: 探查得到的表结构（没设 schema，探查的是当前库）
SCHEMA = {"tables": {"sales": {"qualified": "sales", "schema": None, "columns": [{"name": c} for c in LAYOUT]}}}
ROWS = [[1, "华东", "13900000001", "张伟", 1200, "改址"], [2, "华南", "13900000002", "李娜", 300, "回访"]]
def as_json(row: list) -> str:
return json.dumps(dict(zip(LAYOUT, row)), ensure_ascii=False)
def as_record(row: list) -> str:
return "(" + ",".join(str(v) for v in row) + ")"
def hidden_of(sql: str, columns: list[str], kind: str = "postgres", schema: dict | None = None) -> set[str]:
snapshot = {"source": "shop", "sql": sql, "columns": columns, "rows": []}
return masking._derived_hidden(snapshot, set(MASKED), kind, masking._schema_columns(schema))[0]
def masked(sql: str, columns: list[str], rows: list[list], kind: str = "postgres",
schema: dict | None = SCHEMA) -> masking.QueryMask:
payload = {"columns": columns, "rows": rows, "sql": sql}
return masking.mask_query_result(payload, MASKED, kind, schema=schema)
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, cells", [
pytest.param("SELECT row_to_json(s) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS], id="row_to_json(s)"),
pytest.param("SELECT to_jsonb(s) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS], id="to_jsonb(s)"),
pytest.param("SELECT s AS j, s.id AS n FROM sales s", [as_record(r) for r in ROWS], id="SELECT s"),
pytest.param("SELECT s::text AS j, s.id AS n FROM sales s", [as_record(r) for r in ROWS], id="s::text"),
pytest.param("SELECT to_json(s.*) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS], id="to_json(s.*)"),
pytest.param("SELECT concat_ws(',', s.*) AS j, s.id AS n FROM sales s",
[",".join(str(v) for v in r) for r in ROWS], id="concat_ws(',', s.*)"),
pytest.param("SELECT ROW(s.*) AS j, s.id AS n FROM sales s", [as_record(r) for r in ROWS], id="ROW(s.*)"),
pytest.param("SELECT (s).buyer_phone AS j, s.id AS n FROM sales s", [r[2] for r in ROWS], id="(s).buyer_phone"),
pytest.param("SELECT (s.*).buyer_name AS j, s.id AS n FROM sales s", [r[3] for r in ROWS], id="(s.*).buyer_name"),
pytest.param("SELECT (SELECT row_to_json(s)) AS j, s.id AS n FROM sales s", [as_json(r) for r in ROWS],
id="相关子查询里的整行"),
pytest.param("SELECT kv.value AS j, s.id AS n FROM sales s, LATERAL jsonb_each_text(to_jsonb(s)) AS kv "
"WHERE kv.key = 'buyer_phone'", [r[2] for r in ROWS], id="LATERAL 拆开整行"),
])
def test_a_whole_row_carrying_masked_columns_is_masked(sql, cells):
"""整行、函数参数里的 s.* 把遮罩列的值交给了别的列：这一列遮，编号照常。"""
done = masked(sql, ["j", "n"], [[cell, row[0]] for cell, row in zip(cells, ROWS)])
assert done.columns == ["j"] and not leaks(done), done.payload["rows"]
assert [row[1] for row in done.payload["rows"]] == [1, 2]
@pytest.mark.parametrize("sql", [
"SELECT json_agg(t) AS j FROM (SELECT id, buyer_phone FROM sales) t",
"SELECT row_to_json(t) AS j FROM (SELECT buyer_name AS 名 FROM sales) t",
"WITH t AS (SELECT id, buyer_phone AS p FROM sales) SELECT json_agg(t) AS j FROM t",
])
def test_a_whole_row_of_a_derived_source_is_read_through_its_columns(sql):
"""整行取自子查询、CTE：按它选出的列认，里面有遮罩列（改了名字的也算）就算。"""
assert "j" in hidden_of(sql, ["j"])
@pytest.mark.parametrize("sql", [
"SELECT (SELECT j FROM (SELECT row_to_json(s) AS j) z) AS j FROM sales s",
"SELECT (SELECT j FROM (SELECT json_agg(t) AS j) z) AS j FROM (SELECT id, buyer_phone FROM sales) t",
"SELECT (WITH w AS (SELECT row_to_json(s) AS j) SELECT j FROM w) AS j FROM sales s",
"SELECT (SELECT j FROM (SELECT to_json(s.*) AS j) z) AS j FROM sales s",
])
def test_a_whole_row_of_an_outer_query_is_found_from_a_nested_derived_table(sql):
"""FROM 里的子查询、CTE 看不见同一层 FROM 里的别的来源，却看得见再外层的查询：写在里面的整行引用、s.* 指外层的来源，
照样认得出。"""
assert "j" in hidden_of(sql, ["j"])
@pytest.mark.parametrize("sql", [
"SELECT json_agg(t) AS j FROM (SELECT id, region, amount FROM sales) t",
"SELECT count(s.*) AS j FROM sales s",
"SELECT COUNT(*) AS j FROM sales s",
"SELECT region AS j, SUM(amount) AS total FROM sales s GROUP BY region",
"SELECT EXISTS(SELECT * FROM sales s2 WHERE s2.id = s.id) AS j FROM sales s",
])
def test_rows_and_stars_that_bring_no_masked_values_stay_visible(sql):
"""整行取自没有遮罩列的子查询、count(s.*) 只数个数：不算。"""
assert hidden_of(sql, ["j", "total"]) == set(MASKED)
def test_a_column_named_like_its_table_is_a_column_where_the_schema_says_so():
"""Postgres 先按列认：表里有一列和表同名（visits 表的 visits 列），SELECT visits 取的是这一列，不是整行。表结构里查得到
这一列就不按整行算；SQLite、MySQL 没有整行引用，不看表结构也不按整行算。"""
schema = {"tables": {"visits": {"columns": [{"name": c} for c in ("id", "visits", "buyer_phone")]}}}
sql = "SELECT visits AS j FROM visits"
assert hidden_of(sql, ["j"], schema=schema) == set(MASKED)
assert hidden_of(sql, ["j"], kind="sqlite") == set(MASKED)
assert "j" in hidden_of(sql, ["j"]), "表结构里查不到这一列：可能是整行，按整行算（多遮）"
def test_a_star_inside_a_function_reads_every_source_it_expands():
"""Oracle 的 JSON_OBJECT(*) 把整行拼成 JSON：同 to_json(s.*)。"""
done = masked("SELECT JSON_OBJECT(*) AS j FROM sales", ["j"], [[as_json(r)] for r in ROWS], kind="oracle")
assert done.columns == ["j"] and not leaks(done)
# --------------------------------------------------------------------------
# 列名表、TABLE x
# --------------------------------------------------------------------------
def test_a_table_with_a_column_list_masks_the_renamed_masked_column():
"""数据源里的表带列名表：号码改叫 p。列名表写全了的，按表结构对上位置，只遮 p 和姓名改叫的 q，编号照常；Postgres 的
列名表只写前几列（后面的照原名）时，列名表和表的列数对不上，改了名的几列整段遮（多遮），照原名的姓名按名字遮。"""
full = "SELECT p AS j, i AS n FROM sales AS s(i, r, p, q, a, o)"
done = masked(full, ["j", "n"], [[r[2], r[0]] for r in ROWS])
assert done.columns == ["j"] and not leaks(done)
assert set(masked(full, ["j", "n"], [[r[2], r[0]] for r in ROWS], schema=None).columns) == {"j", "n"}, \
"表结构里查不到这张表：说不清哪几列改了名，都遮"
short = masked("SELECT p AS j, i AS n FROM sales AS s(i, r, p)", ["j", "n"], [[r[2], r[0]] for r in ROWS])
assert set(short.columns) == {"j", "n"} and not leaks(short)
star = masked("SELECT * FROM sales AS s(i, r, p)", ["i", "r", "p", *LAYOUT[3:]], [list(r) for r in ROWS])
assert star.columns == ["i", "r", "p", "buyer_name"] and not leaks(star)
def test_a_column_list_on_a_cte_reference_renames_it_too():
"""引用 CTE 时再写列名表（FROM t AS x(p, a)）：同样按位置对上。"""
sql = "WITH t AS (SELECT buyer_phone, amount FROM sales) SELECT p AS j, a AS total FROM t AS x(p, a)"
assert hidden_of(sql, ["j", "total"]) >= {"j"}
assert "total" not in hidden_of(sql, ["j", "total"])
@pytest.mark.parametrize("sql, columns, shown", [
pytest.param("WITH t(a, b, c, d, e, f) AS (TABLE sales) SELECT c AS j, a AS n FROM t", ["j", "n"], {"n"},
id="CTE 里的 TABLE"),
pytest.param("SELECT * FROM (TABLE sales) AS t(a, b, c, d, e, f)", ["a", "b", "c", "d", "e", "f"], {"a", "b", "e"},
id="子查询里的 TABLE"),
pytest.param("WITH t(a, b, c) AS (TABLE sales) SELECT c AS j FROM t", ["j"], set(), id="列名表只写前几列"),
])
def test_a_table_statement_is_read_as_select_star(sql, columns, shown):
"""TABLE sales（Postgres、MySQL）就是 SELECT * FROM sales：列名表照样按位置对上，号码改叫 c 的那一列遮。"""
hidden = hidden_of(sql, columns, schema=SCHEMA)
assert {"j", "c"} & hidden and not shown & hidden, hidden
@pytest.mark.parametrize("sql", [
"SELECT 0 AS a, '' AS b, '' AS c, '' AS d, 0 AS e, '' AS f WHERE false UNION ALL TABLE sales",
"WITH t(a, b, c) AS (TABLE archive.sales) SELECT c AS j FROM t",
"WITH t(a, b, c) AS (TABLE sales ORDER BY id LIMIT 1) SELECT c AS j FROM t",
"WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT c AS j FROM t WHERE (",
])
def test_sql_that_cannot_be_read_but_has_a_star_or_a_table_statement_hides_everything(sql):
"""解析器认不出的写法（UNION 后一支的 TABLE、带库名或 ORDER BY 的 TABLE）：SQL 里没写遮罩列，可 TABLE 语句、* 会把它
带进别的列名，整份快照的列都算，同解析不了又写着遮罩列的。"""
columns = ["a", "b", "c", "j"]
assert hidden_of(sql, columns) >= set(columns)
@pytest.mark.parametrize("sql", [
"SELECT COUNT(*) AS j FROM sales WHERE (",
"SELECT amount * 2 AS j FROM sales WHERE (",
])
def test_sql_that_cannot_be_read_with_only_a_count_or_a_product_shows_the_rest(sql):
"""只有 COUNT(*)、乘号：不会把遮罩列带出来，照旧只遮遮罩列本身。"""
assert hidden_of(sql, ["j"]) == set(MASKED)
# --------------------------------------------------------------------------
# UNPIVOT、PIVOT、MATCH_RECOGNIZE
# --------------------------------------------------------------------------
@pytest.mark.parametrize("sql, columns", [
pytest.param("SELECT id, val FROM sales UNPIVOT (val FOR col IN (buyer_phone, region))", ["id", "val"],
id="UNPIVOT"),
pytest.param("SELECT * FROM (SELECT region, buyer_phone FROM sales) PIVOT (MAX(buyer_phone) FOR region IN "
"('华东' AS e, '华南' AS s))", ["e", "s"], id="PIVOT"),
pytest.param("SELECT * FROM sales MATCH_RECOGNIZE (PARTITION BY region ORDER BY id MEASURES buyer_phone AS x "
"ONE ROW PER MATCH PATTERN (a) DEFINE a AS amount > 0)", ["region", "x"], id="MATCH_RECOGNIZE"),
])
def test_pivots_reading_masked_columns_hide_everything(sql, columns):
"""UNPIVOT 的值列、PIVOT 按 IN 里的值起的列、MATCH_RECOGNIZE 的 MEASURES 装着遮罩列的值：出的是哪几列不好认，整份遮。"""
assert hidden_of(sql, columns, kind="oracle") >= {c.lower() for c in columns}
@pytest.mark.parametrize("sql, columns", [
("SELECT id, val FROM sales UNPIVOT (val FOR col IN (region, note))", ["id", "val"]),
("SELECT * FROM (SELECT region, buyer_phone FROM sales) PIVOT (COUNT(buyer_phone) FOR region IN "
"('华东' AS e, '华南' AS s))", ["e", "s"]),
])
def test_pivots_that_bring_no_masked_values_stay_visible(sql, columns):
assert hidden_of(sql, columns, kind="oracle") == set(MASKED)
# --------------------------------------------------------------------------
# 选出项套了括号的 *
# --------------------------------------------------------------------------
def test_a_star_in_parentheses_is_counted_as_a_star():
"""Postgres 里选出项 (s.*)、(s).* 照样展开成好几列：后面没起别名的表达式按结果末尾倒着数，不按一列数位置。"""
columns = [*LAYOUT, "upper"]
for item in ("(s.*)", "(s).*"):
sql = f"SELECT {item}, UPPER(s.buyer_name) FROM sales s"
assert hidden_of(sql, columns) == {*MASKED, "upper"}, item
def test_sqlite_runs_of_the_same_shapes_are_unchanged():
"""SQLite 没有整行引用、列名表：同样的写法照旧（表和列同名时按列认，不多遮）。"""
conn = sqlite3.connect(":memory:")
conn.executescript("CREATE TABLE region (id INTEGER, region TEXT, buyer_phone TEXT);"
"INSERT INTO region VALUES (1, '华东', '13900000001');")
sql = "SELECT region AS 区域, COUNT(*) AS n FROM region GROUP BY region"
cursor = conn.execute(sql)
payload = {"columns": [d[0] for d in cursor.description], "rows": [list(r) for r in cursor.fetchall()], "sql": sql}
assert masking.mask_query_result(payload, MASKED, "sqlite").payload is payload
conn.close()
backend/tests/test_verify_agent_masking.py
+625
View File
@@ -0,0 +1,625 @@
"""核查 agent 看到的查询结果先遮、存下的快照也遮（engine/verify_agent.py 建查询工具时打开 tools/datasource 的遮罩开关）。
双盲复算、反方审稿、审批质询共用的核查 agent 以前把查询结果原样交给模型，遮罩是事后按值做的（审批质询的 _Secrets），
每一轮复核都找到新的漏法。现在在工具边界先遮：数据源设了遮罩时，按 SQL 认出取自遮罩列的结果列（别名、函数、子查询、
WITH、SELECT *），连同列名直接对上的，这些格子在交给模型的结果、落库的快照里都是「已遮罩」，快照记下哪些列在查询时
就遮了；别的文字格里夹着这份结果里遮罩值的那一截同样换掉（换了写法的号码也认）。运行详情「工件」页签走的通用接口
GET /api/artifacts/{id} 给的也就是遮好的。没碰遮罩列的聚合查询原样。遮掉的值只在内存里交回调用方（审批质询按值遮的
第二道防线用），不进快照、不交给模型。遮罩开关缺省关着，别的调用方（Agent 节点、调用工具节点）照旧。
模型一律是 MockChatModel 剧本，数据源是现造的 SQLite 文件，手机号、姓名、金额都是虚构的。
"""
from __future__ import annotations
import json
import sqlite3
import uuid
import pytest
from httpx import ASGITransport, AsyncClient
from sqlalchemy import delete
from app.core import artifact_store
from app.core.events import EventType
from app.data import introspect
from app.data.engine import engines
from app.db.base import SessionLocal
from app.db.models import DataSource, Setting
from app.engine import masking, spend
from app.engine.verify_agent import run_verify_agent
from app.main import app
from app.tools.datasource import build_datasource_tools
from app.tools.registry import ToolContext, Versions
from tests.test_verify_agent import SPEC, Script, q
MASKED = "已遮罩"
#: 遮罩列里的值（和备注里换了写法的号码）：交给模型的、存下的、接口给的都不能出现
SECRETS = ("1390000000", "张伟", "李娜", "王芳", "刘洋", "139-0000-0003", "139 0000 0003")
SALES = ("CREATE TABLE sales (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, buyer_name TEXT, amount INTEGER,"
" note TEXT);"
"INSERT INTO sales VALUES (1, '华东', '13900000001', '张伟', 1200, '客户13900000001来电改址'),"
" (2, '华东', '13900000002', '李娜', 300, '例行回访'), (3, '华南', '13900000003', '王芳', 900, '回访：联系 139-0000-0003'),"
" (4, '华北', '13900000004', '刘洋', 400, '');"
# 联系方式表的号码列、客户资料的 JSON 列都没设遮罩（json_shop 另设了 profile）：同一批号码、姓名写在这里
"CREATE TABLE contacts (id INTEGER PRIMARY KEY, phone_num TEXT, visits INTEGER);"
"INSERT INTO contacts VALUES (1, '13900000001', 5), (2, '139 0000 0003', 2), (3, '13800000000', 1);"
"CREATE TABLE customers (id INTEGER PRIMARY KEY, region TEXT, profile TEXT);"
"INSERT INTO customers VALUES (1, '华东', '{\"phone\": \"13900000001\", \"name\": \"张伟\"}'),"
" (2, '华南', '{\"phone\": \"13900000002\", \"name\": \"李娜\"}');")
VALUE_SCHEMA = {"type": "object", "properties": {"value": {"type": "number", "description": "要核对的数"}},
"required": ["value"]}
@pytest.fixture(autouse=True)
async def clean_spend():
async with SessionLocal() as session:
await session.execute(delete(Setting).where(Setting.key.in_([spend.SPEND_KEY, "verify"])))
await session.commit()
yield
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(script)
db.commit()
db.close()
name = f"shop_{uuid.uuid4().hex[:8]}"
async with SessionLocal() as session:
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=True, enabled=True, options=options or {})
session.add(row)
await session.flush()
row.schema_cache = await introspect.introspect(row)
await session.commit()
return name, row.id
@pytest.fixture
async def masked_shop(tmp_path):
"""四笔订单，买家手机号、姓名两列设了遮罩；备注没设遮罩，其中两条写着手机号（一条换了写法）。返回源名。"""
name, source_id = await _source(tmp_path, {"mask_columns": "buyer_phone, buyer_name"})
yield name
await engines.invalidate(source_id)
@pytest.fixture
async def json_shop(tmp_path):
"""同一个库，客户资料（JSON：手机号、姓名）这一列设了遮罩。返回源名。"""
name, source_id = await _source(tmp_path, {"mask_columns": "profile"})
yield name
await engines.invalidate(source_id)
@pytest.fixture
async def plain_shop(tmp_path):
name, source_id = await _source(tmp_path, None)
yield name
await engines.invalidate(source_id)
async def ask(name: str, **kwargs):
"""按剧本跑一次核查 agent（剧本由各个测试先用 Script 装好），交回结果和发出的事件。"""
events: list[tuple] = []
out = await run_verify_agent(spec=SPEC, system="核对一个数。", task="查一下。", sources=[name],
output_schema=VALUE_SCHEMA, data_versions=Versions.CURRENT, run_id="run-mask",
node_id="check", emit=lambda kind, **data: events.append((kind, data)), **kwargs)
return out, events
def _clean(text: str) -> bool:
return not any(s in text for s in SECRETS)
def _cells(snapshot: dict, column: str) -> list:
i = snapshot["columns"].index(column)
return [row[i] for row in snapshot["rows"]]
@pytest.mark.parametrize("sql, masked", [
pytest.param("SELECT buyer_phone, amount FROM sales ORDER BY id", ["buyer_phone"], id="直接选出"),
pytest.param("SELECT buyer_phone AS 手机号, amount AS 金额 FROM sales ORDER BY id", ["手机号"], id="起别名"),
pytest.param("SELECT UPPER(buyer_name) AS 姓名, amount FROM sales ORDER BY id", ["姓名"], id="套函数"),
pytest.param("SELECT t.p AS 号码, t.amount FROM (SELECT id, buyer_phone AS p, amount FROM sales) t ORDER BY t.id",
["号码"], id="子查询"),
pytest.param("WITH t AS (SELECT id, buyer_phone AS p, amount FROM sales) SELECT p, amount FROM t ORDER BY id",
["p"], id="WITH"),
pytest.param("SELECT * FROM sales ORDER BY id", ["buyer_phone", "buyer_name"], id="SELECT *"),
])
async def test_masked_columns_are_masked_before_the_model_sees_them(monkeypatch, masked_shop, sql, masked):
name = masked_shop
amount = "amount" if "金额" not in sql else "金额"
script = Script(monkeypatch, [q(name, sql, "c1")],
{"value": {"alias": "Q1", "row": 0, "column": amount}, "rationale": "见 [[see:Q1]]。"})
out, events = await ask(name)
assert out.error is None and out.result["value"] == 1200, "没遮的列照常取值"
# 交给模型的工具结果：遮罩列是「已遮罩」，别的列照原样
seen = script.tool_results()["c1"]
assert seen.startswith("【证据 Q1】") and MASKED in seen and _clean(seen), seen
payload = json.loads(seen.split("\n", 1)[1])
for column in masked:
assert set(_cells(payload, column)) == {MASKED}
assert _cells(payload, amount) == [1200, 300, 900, 400]
assert "mask_columns" not in payload and "masked_at_query" not in payload, "交给模型的结果不多出字段"
# 落库的快照：同样遮好，记下查询时就遮了的列；数据源设的遮罩列照记，换了名字的结果列也记进去（证据面板据此画成遮罩）
[query] = out.queries
snapshot = artifact_store.load(query["artifact"])
assert _clean(json.dumps(snapshot, ensure_ascii=False))
for column in masked:
assert set(_cells(snapshot, column)) == {MASKED}
assert snapshot.get("column_types", {}).get(column, "text") == "text"
assert _cells(snapshot, amount) == [1200, 300, 900, 400]
assert snapshot["masked_at_query"] == masked
# mask_columns 照记查询当时数据源设的遮罩列（证据面板按它从表的字段清单里去掉列），改了名字的结果列只记在
# masked_at_query：混进去的话，同源别的表里同名的真实字段（name、region）会从字段清单里消失
assert snapshot["mask_columns"] == ["buyer_phone", "buyer_name"]
# 运行详情「工件」页签走的通用接口：给的就是遮好的
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
got = await client.get(f"/api/artifacts/{query['artifact']}")
assert got.status_code == 200 and got.json()["content"] == snapshot and _clean(got.text)
# 事件（tool.end 的预览）、收尾时交给模型的查询结果里也没有
[end] = [d for k, d in events if k == EventType.TOOL_END]
assert MASKED in end["preview"] and _clean(end["preview"])
final = "\n".join(str(m.content) for m in script.calls[-1][1])
assert MASKED in final and _clean(final)
# 遮掉的值只在内存里交回调用方，不跟着结果往外交
assert out.masked_values is not None and out.masked_values.strong
assert "masked_values" not in out.as_dict()
@pytest.mark.parametrize("sql, value", [
("SELECT SUM(amount) AS total FROM sales", 2800),
("SELECT COUNT(DISTINCT buyer_phone) AS total FROM sales", 4),
("SELECT SUM(amount) AS total FROM sales WHERE buyer_phone = '13900000001'", 1200),
("SELECT region, SUM(amount) AS total FROM sales GROUP BY region ORDER BY total DESC", 1500),
("SELECT region, note, amount AS total FROM sales ORDER BY id", 1200),
# 汇总里的 CASE 只在条件里用到手机号：同写在 WHERE 里，值全部来自金额
("SELECT SUM(CASE WHEN buyer_phone LIKE '139%' THEN amount ELSE 0 END) AS total FROM sales", 2800),
])
async def test_queries_that_do_not_read_masked_values_are_untouched(monkeypatch, masked_shop, sql, value):
"""没碰遮罩列的值（汇总、只数个数、只在条件里用到）：结果原样，取值照常，快照记下查询时一列都没遮。只遮这份结果
里的遮罩值：备注里写着手机号、这份结果却没有选出手机号的，按值认不出（那是数据源的遮罩没设全）。"""
name = masked_shop
script = Script(monkeypatch, [q(name, sql, "c1")],
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "见 [[see:Q1]]。"})
out, _ = await ask(name)
assert out.error is None and out.result["value"] == value
seen = script.tool_results()["c1"]
assert MASKED not in seen
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert snapshot["masked_at_query"] == [] and MASKED not in json.dumps(snapshot, ensure_ascii=False)
assert [c.lower() for c in snapshot["mask_columns"]] == ["buyer_phone", "buyer_name"]
async def test_masked_values_inside_other_text_cells_are_cut_out(monkeypatch, masked_shop):
"""同一份结果里选出了手机号：备注里夹着的同一个号码（换了写法的也算）换成「已遮罩」，备注别的字照原样。"""
name = masked_shop
sql = "SELECT buyer_phone, note, amount AS total FROM sales ORDER BY id"
script = Script(monkeypatch, [q(name, sql, "c1")],
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "见 [[see:Q1]]。"})
out, _ = await ask(name)
seen = script.tool_results()["c1"]
assert _clean(seen)
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert _cells(snapshot, "note") == [f"客户{MASKED}来电改址", "例行回访", f"回访：联系 {MASKED}", ""]
assert snapshot["masked_at_query"] == ["buyer_phone"], "备注只换了一截，不算整列遮"
async def test_a_cell_reference_into_a_masked_column_is_not_taken(monkeypatch, masked_shop):
"""模型交的格子引用指到查询时就遮了的那一列：不采信，照实说是遮罩列，不说成「不是数值」。"""
name = masked_shop
sql = "SELECT buyer_phone AS 手机号, amount FROM sales ORDER BY id"
Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "手机号"}, "rationale": "x"})
out, _ = await ask(name)
assert out.result["value"] is None
[cell] = out.cells
assert cell["status"] == "unresolved" and "1390000000" not in json.dumps(cell)
assert cell["reason"] == "Q1 的「手机号」取自设置了遮罩的列，查询时已遮罩，不采信"
async def test_sql_that_cannot_be_read_but_names_a_masked_column_masks_the_whole_result(monkeypatch, masked_shop):
"""SQL 认不出哪一列取自遮罩列、里面又写着遮罩列：整份结果的列都遮（和双盲复算、审批质询同一个保守规则）。"""
name = masked_shop
monkeypatch.setattr(masking, "_tainted", lambda *args, **kwargs: None)
sql = "SELECT buyer_phone, amount FROM sales ORDER BY id"
Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "amount"}, "rationale": "x"})
out, _ = await ask(name)
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert snapshot["masked_at_query"] == ["buyer_phone", "amount"]
assert set(_cells(snapshot, "amount")) == {MASKED} and out.result["value"] is None
async def test_the_switch_is_off_by_default_and_a_source_without_masks_is_untouched(masked_shop, plain_shop):
"""遮罩开关缺省关着：Agent 节点、调用工具节点建的查询工具照旧交回原样的结果。没设遮罩的数据源开着开关也一字不差
（快照里不多出键）。"""
sql = "SELECT buyer_phone, amount FROM sales ORDER BY id"
ctx = ToolContext(run_id="run-switch", node_id="n1", data_versions=Versions.CURRENT)
async with SessionLocal() as session:
[plain_tool] = await build_datasource_tools([f"db_query__{masked_shop}"], ctx, session)
[masking_tool] = await build_datasource_tools([f"db_query__{masked_shop}"], ctx, session, mask_results=True)
[unmasked_tool] = await build_datasource_tools([f"db_query__{plain_shop}"], ctx, session, mask_results=True)
raw = json.loads(await plain_tool.ainvoke({"sql": sql}))
assert _cells(raw, "buyer_phone")[0] == "13900000001"
stored = artifact_store.load(raw["artifact"])
assert _cells(stored, "buyer_phone")[0] == "13900000001" and "masked_at_query" not in stored
shown = json.loads(await masking_tool.ainvoke({"sql": sql}))
assert set(_cells(shown, "buyer_phone")) == {MASKED}
untouched = json.loads(await unmasked_tool.ainvoke({"sql": sql}))
stored = artifact_store.load(untouched["artifact"])
assert _cells(stored, "buyer_phone")[0] == "13900000001"
assert "masked_at_query" not in stored and "mask_columns" not in stored
async def test_column_names_the_model_wrote_are_kept(monkeypatch, masked_shop):
"""列名是模型自己写在 SQL 里的（哪怕写着一个号码）：工具边界只遮格子、不改列名，它引用的列名才对得上。"""
name = masked_shop
sql = 'SELECT buyer_phone, amount AS "13900000001 的金额" FROM sales ORDER BY id'
Script(monkeypatch, [q(name, sql, "c1")],
{"value": {"alias": "Q1", "row": 0, "column": "13900000001 的金额"}, "rationale": "x"})
out, _ = await ask(name)
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert snapshot["columns"] == ["buyer_phone", "13900000001 的金额"]
assert set(_cells(snapshot, "buyer_phone")) == {MASKED} and out.result["value"] == 1200
# --------------------------------------------------------------------------
# 别的查询里遮掉的值、调用方给的值：后面的查询按值遮（复核第 1 轮）
# --------------------------------------------------------------------------
async def test_values_masked_in_an_earlier_query_are_masked_in_later_ones(monkeypatch, masked_shop):
"""Q1 选出手机号（查询时整列遮掉，值并进内存）；之后的查询里，没设遮罩的联系方式表、备注写着同一批号码（换了写法的
也算）：整格是号码的整格换，夹在文字里的换那一截。交给模型的、落库的、通用接口给的都是遮好的。以前只按这一份结果
里的遮罩值换，Q2 原样交给模型、原样落库，质询工件事后也收不到 Q1 的值。"""
name = masked_shop
steps = [q(name, "SELECT buyer_phone, amount AS total FROM sales ORDER BY id", "c1"),
q(name, "SELECT phone_num, visits FROM contacts ORDER BY id", "c2"),
q(name, "SELECT id, note FROM sales ORDER BY id", "c3")]
script = Script(monkeypatch, steps, {"value": {"alias": "Q2", "row": 0, "column": "visits"}, "rationale": "x"})
out, _ = await ask(name)
assert out.error is None and out.result["value"] == 5
seen = script.tool_results()
assert all(_clean(seen[c]) for c in ("c1", "c2", "c3")), seen
contacts = artifact_store.load(out.queries[1]["artifact"])
assert contacts["rows"] == [[MASKED, 5], [MASKED, 2], ["13800000000", 1]]
assert contacts["masked_at_query"] == [], "按值换的不算整列遮：没对上的号码照常显示"
notes = artifact_store.load(out.queries[2]["artifact"])
assert _cells(notes, "note") == [f"客户{MASKED}来电改址", "例行回访", f"回访：联系 {MASKED}", ""]
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
for query in out.queries:
got = await client.get(f"/api/artifacts/{query['artifact']}")
assert got.status_code == 200 and _clean(got.text)
async def test_values_given_by_the_caller_are_masked_before_the_model_sees_them(monkeypatch, masked_shop):
"""调用方交进来的遮罩值（审批质询、反方审稿收的报告快照里的值）：第一条查询就按它遮，模型看不到。"""
name = masked_shop
given = masking._Secrets()
given.add({"columns": ["buyer_phone"], "rows": [["13900000001"], ["13900000003"]]}, {"buyer_phone"})
script = Script(monkeypatch, [q(name, "SELECT id, note, amount AS total FROM sales ORDER BY id", "c1")],
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "x"})
out, _ = await ask(name, masked_values=given)
assert out.result["value"] == 1200 and _clean(script.tool_results()["c1"])
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert _cells(snapshot, "note") == [f"客户{MASKED}来电改址", "例行回访", f"回访：联系 {MASKED}", ""]
assert not given.strong - out.masked_values.strong, "交进来的值随遮掉的值一起交回"
async def test_a_reference_to_a_cell_masked_by_value_says_so(monkeypatch, masked_shop):
"""格子引用指到按值换成「已遮罩」的那一格：照实说，不说成「不是数值」。"""
name = masked_shop
steps = [q(name, "SELECT buyer_phone FROM sales ORDER BY id", "c1"),
q(name, "SELECT phone_num, visits FROM contacts ORDER BY id", "c2")]
Script(monkeypatch, steps, {"value": {"alias": "Q2", "row": 0, "column": "phone_num"}, "rationale": "x"})
out, _ = await ask(name)
assert out.result["value"] is None
cell = out.cells[0]
assert cell["status"] == "unresolved"
assert cell["reason"] == "Q2 第 0 行「phone_num」与设置了遮罩的值相同，查询时已遮罩，不采信"
async def test_a_table_function_over_a_masked_json_column_is_masked(monkeypatch, json_shop):
"""遮罩的是 JSON 列（客户资料）：用 json_each 拆开是最自然的写法，拆出来的 key、value 在交给模型、落库之前就遮掉。
以前原样交给模型并落库，masked_at_query 是空的。"""
name = json_shop
sql = "SELECT c.id AS 编号, j.key, j.value FROM customers c, json_each(c.profile) j ORDER BY c.id, j.key"
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"] == ["key", "value"] and _cells(snapshot, "编号") == [1, 1, 2, 2]
assert _clean(json.dumps(snapshot, ensure_ascii=False))
async def test_a_failed_query_does_not_hand_masked_values_to_the_model(monkeypatch, masked_shop):
"""数据库的报错会带上数据里的值（Postgres 类型转换失败：invalid input syntax for type bigint: "139-0000-0001"）：
数据源设了遮罩时，引号里、SQL 里没写着的那一段和不带引号的长数换成「已遮罩」；模型自己写的名字照留（它据此改写
SQL）。遮罩开关没开的（Agent 节点）照旧原话。"""
from app.tools import datasource
said = {"message": ""}
async def failing(source, sql, **kwargs):
raise RuntimeError(said["message"])
monkeypatch.setattr(datasource, "run_query", failing)
sql = "SELECT CAST(buyer_phon AS BIGINT) AS p FROM sales WHERE id = 12345"
ctx = ToolContext(run_id="run-error", node_id="n1", data_versions=Versions.CURRENT)
async with SessionLocal() as session:
[masking_tool] = await build_datasource_tools([f"db_query__{masked_shop}"], ctx, session, mask_results=True)
[plain_tool] = await build_datasource_tools([f"db_query__{masked_shop}"], ctx, session)
for message, shown in [
('invalid input syntax for type bigint: "139-0000-0001"', 'invalid input syntax for type bigint: "已遮罩"'),
("value '13900000001' is out of range for type integer", "value '已遮罩' is out of range for type integer"),
("Conversion Error: Type INT64 with value 13900000001 can't be cast to INT32",
"Conversion Error: Type INT64 with value 已遮罩 can't be cast to INT32"),
('column "buyer_phon" does not exist', 'column "buyer_phon" does not exist'),
("no rows near id 12345", "no rows near id 12345"),
]:
said["message"] = message
assert await masking_tool.ainvoke({"sql": sql}) == f"查询失败：{shown}"
assert await plain_tool.ainvoke({"sql": sql}) == f"查询失败：{message}"
def test_a_result_from_a_source_without_masks_is_masked_by_the_values_seen_so_far():
"""数据源没设遮罩（另一个库里的联系方式表），此前查询遮掉的值照样按值换；这一份结果没有可收的值。"""
known = masking._Secrets()
known.add({"columns": ["buyer_phone"], "rows": [["13900000001"]]}, {"buyer_phone"})
payload = {"sql": "SELECT phone_num, note FROM contacts", "columns": ["phone_num", "note"],
"rows": [["13900000001", "客户 139 0000 0001 来电"], ["13800000000", "例行回访"]]}
done = masking.mask_query_result(payload, [], "sqlite", known)
assert done.payload["rows"] == [[MASKED, f"客户 {MASKED} 来电"], ["13800000000", "例行回访"]]
assert done.columns == [] and not done.secrets.strong
assert masking.mask_query_result(payload, [], "sqlite").payload is payload, "什么都没收到时原样"
def test_labels_written_in_this_query_are_not_masked_by_values_from_earlier_ones():
"""分档标签是这条 SQL 自己写的字：此前查询里认不准、当遮罩值收了的同名标签，不把这一份的标签列换掉。"""
known = masking._Secrets()
known.add({"columns": ["档"], "rows": [["大单"], ["小单"]]}, {"档"})
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)
# --------------------------------------------------------------------------
# 过期的表结构、写了库名的表、整行与列名表（复核第 3 轮）
# --------------------------------------------------------------------------
#: 一月那张表后来多了一列 extra（排在第二），数据源的表结构是那之前探查的；二月那张号码排在第二列
REBUILT = ("CREATE TABLE jan (id INTEGER PRIMARY KEY, extra TEXT, region TEXT, buyer_phone TEXT, amount INTEGER);"
"CREATE TABLE feb (id INTEGER PRIMARY KEY, buyer_phone TEXT, region TEXT, amount INTEGER, extra TEXT);"
"INSERT INTO jan VALUES (1, '加急', '华东', '13900000001', 1200);"
"INSERT INTO feb VALUES (2, '13900000002', '华南', 300, '普通');")
async def test_a_stale_schema_does_not_place_masked_values_by_its_old_layout(monkeypatch, tmp_path):
"""数据源的表结构过期了（一月那张少了 extra）：UNION 第一支的 * 按它数出的列名和后一支的列数对不上，不按它对位置，整份
遮。以前照过期的列名对，二月的号码落进的「extra」认不出，原样交给模型、落库。"""
name, source_id = await _source(tmp_path, {"mask_columns": "buyer_phone"}, script=REBUILT)
async with SessionLocal() as session:
row = await session.get(DataSource, source_id)
cache = dict(row.schema_cache)
jan = dict(cache["tables"]["jan"])
jan["columns"] = [c for c in jan["columns"] if c["name"] != "extra"]
row.schema_cache = {**cache, "tables": {**cache["tables"], "jan": jan}}
await session.commit()
try:
sql = ("WITH t AS (SELECT * FROM jan UNION ALL SELECT id, buyer_phone, region, amount, extra FROM feb) "
"SELECT extra AS 备注, id AS n FROM t ORDER BY id")
script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "n"},
"rationale": "x"})
out, _ = await ask(name)
seen = script.tool_results()["c1"]
assert MASKED in seen and _clean(seen), seen
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert "备注" in snapshot["masked_at_query"] and _clean(json.dumps(snapshot, ensure_ascii=False))
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 _clean(got.text)
finally:
await engines.invalidate(source_id)
#: 造出来的 Postgres、Oracle 数据源的表结构：探查的是当前库（没设 schema），sales 一张表
PG_LAYOUT = ["id", "region", "buyer_phone", "buyer_name", "amount", "note"]
PG_SCHEMA = {"tables": {"sales": {"qualified": "sales", "schema": None,
"columns": [{"name": c, "type": "TEXT"} for c in PG_LAYOUT]}}, "schema": None}
PG_ROWS = [[1, "华东", "13900000001", "张伟", 1200, "改址"], [2, "华南", "13900000002", "李娜", 300, "回访"]]
def _json_row(row: list) -> str:
return json.dumps(dict(zip(PG_LAYOUT, row)), ensure_ascii=False)
#: (数据源类型, SQL, 数据库交回的行, 查询时就整列遮掉的列, 编号那一格取到的值)。结果列按 SQL 里起的别名（j、n）。行照
#: 驱动交回的样子造：JSON 是字符串，记录转成文字；archive.sales 是另一个库里的同名表，号码排在第二列。表的列名表只写了
#: 前几列（Postgres 后面的列照原名）、另一个库里的同名表：说不清号码改叫了哪个名字，那一段整段遮，编号也遮
DRIVER_CASES = [
("postgresql", "SELECT json_agg(t) AS j, 1 AS n FROM (SELECT id, buyer_phone FROM sales) t",
[[json.dumps([{"id": r[0], "buyer_phone": r[2]} for r in PG_ROWS]), 1]], ["j"], 1),
("postgresql", "SELECT row_to_json(s) AS j, s.id AS n FROM sales s ORDER BY s.id",
[[_json_row(r), r[0]] for r in PG_ROWS], ["j"], 1),
("postgresql", "SELECT to_json(s.*) AS j, s.id AS n FROM sales s ORDER BY s.id",
[[_json_row(r), r[0]] for r in PG_ROWS], ["j"], 1),
("postgresql", "SELECT s AS j, s.id AS n FROM sales s ORDER BY s.id",
[["(" + ",".join(map(str, r)) + ")", r[0]] for r in PG_ROWS], ["j"], 1),
("postgresql", "SELECT p AS j, i AS n FROM sales AS s(i, r, p) ORDER BY i", [[r[2], r[0]] for r in PG_ROWS],
["j", "n"], None),
("postgresql", "WITH t(a, b, c, d, e, f) AS (TABLE sales) SELECT c AS j, a AS n FROM t ORDER BY a",
[[r[2], r[0]] for r in PG_ROWS], ["j"], 1),
("postgresql", "WITH t(a, b, c, d, e, f) AS (SELECT * FROM archive.sales) SELECT b AS j, a AS n FROM t",
[["13900000007", 7]], ["j", "n"], None),
("postgresql", "WITH t(a, b, c, d, e, f) AS (SELECT * FROM sales) SELECT c AS j, a AS n FROM t ORDER BY a",
[[r[2], r[0]] for r in PG_ROWS], ["j"], 1),
("oracle", "SELECT JSON_OBJECT(*) AS j, id AS n FROM sales ORDER BY id", [[_json_row(r), r[0]] for r in PG_ROWS],
["j"], 1),
("oracle", "SELECT id AS n, val AS j FROM sales UNPIVOT (val FOR col IN (buyer_phone, region)) ORDER BY id",
[[1, "13900000001"], [1, "华东"], [2, "13900000002"], [2, "华南"]], ["n", "j"], None),
]
@pytest.fixture
async def driver_shop():
"""造一个 Postgres / Oracle 数据源（查询交给下面换掉的 run_query，不连库）：买家手机号、姓名设了遮罩。返回建源的函数。"""
made: list[int] = []
async def make(kind: str) -> str:
name = f"pg_{uuid.uuid4().hex[:8]}"
async with SessionLocal() as session:
row = DataSource(name=name, kind=kind, host="127.0.0.1", port=5432, username="probe", database="shop",
readonly=True, enabled=True, options={"mask_columns": "buyer_phone, buyer_name"},
schema_cache=PG_SCHEMA)
session.add(row)
await session.commit()
made.append(row.id)
return name
yield make
for source_id in made:
await engines.invalidate(source_id)
@pytest.mark.parametrize("kind, sql, rows, masked, value", DRIVER_CASES,
ids=["json_agg(t)", "row_to_json(s)", "to_json(s.*)", "SELECT s", "表的列名表", "CTE 里的 TABLE",
"另一个库里的同名表", "本库的表按位置数", "JSON_OBJECT(*)", "UNPIVOT"])
async def test_whole_rows_and_renamed_columns_from_other_databases_are_masked(monkeypatch, driver_shop, kind, sql,
rows, masked, value):
"""Postgres 的整行（json_agg(t)、row_to_json(s)、to_json(s.*)、SELECT s）、表的列名表、CTE 里的 TABLE sales、另一个库里的
同名表，Oracle 的 JSON_OBJECT(*)、UNPIVOT：交给模型的、落库的、通用接口给的都是遮好的，编号照常（认不出位置的整份遮）。
以前这些写法 masked_at_query 是空的，号码、姓名原样交给模型、落库。"""
from app.data.engine import QueryResult
from app.tools import datasource
async def driver(source, statement, **kwargs):
assert statement == sql
columns = ["n", "j"] if statement.startswith("SELECT id AS n") else ["j", "n"]
return QueryResult(columns=columns, rows=[list(r) for r in rows], row_count=len(rows), truncated=False,
elapsed_ms=1, sql=statement)
monkeypatch.setattr(datasource, "run_query", driver)
name = await driver_shop(kind)
script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "n"},
"rationale": "x"})
out, events = await ask(name)
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 _clean(json.dumps(snapshot, ensure_ascii=False))
assert out.result["value"] == value
[end] = [d for k, d in events if k == EventType.TOOL_END]
assert _clean(end["preview"])
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 _clean(got.text)