feat(pii): self-service per-column masking policy enforced for the LLM

- pii_catalog: persistent per-column masking policy (default masked); GET/POST
  /api/pii/policy and POST /api/pii/lookup which redacts masked values server-side.
- get_pii masked flag now reflects the policy; dataflow exposes the dataset key.
- Data Flow PII inspector: per-column lock/unlock toggles + mask-all/unmask-all,
  so operators control exactly which data the assistant may reveal.

api/dataflow.py

@@ -143,6 +143,7 @@ def _build() -> dict[str, Any]:
         p = pii_by_node.get(n["id"])
         if p:
             node["pii"] = {
+                "key": p["key"],
                 "has_pii": p["has_pii"], "pii_count": p["pii_count"], "all_masked": p["all_masked"],
                 "masked_layer": p.get("masked_layer", False),
                 "categories": sorted({c["category"] for c in p["pii_columns"]}),

api/pii_catalog.py

@@ -8,14 +8,17 @@ per dataset which columns are PII and whether they are masked.
 
 from __future__ import annotations
 
+import json
 import os
 import re
 import time
+from pathlib import Path
 from typing import Any
 
 import httpx
 from fastapi import APIRouter
 from fastapi.responses import JSONResponse
+from pydantic import BaseModel
 
 router = APIRouter(prefix="/api/pii", tags=["pii"])
 
@@ -24,6 +27,14 @@ TRINO_USER = os.getenv("TRINO_USER", "mo")
 OPENMETADATA_URL = os.getenv("OPENMETADATA_URL", "").rstrip("/")
 OPENMETADATA_TOKEN = os.getenv("OPENMETADATA_TOKEN", "")
 
+# User-controlled masking policy. Each PII column can be masked (hidden from the
+# LLM) or unmasked. Persisted to disk so it survives restarts. The default is
+# "masked" — sensitive data is withheld until an operator explicitly opts out.
+POLICY_PATH = Path(os.getenv("MASKING_POLICY_PATH", "/data/masking_policy.json"))
+DEFAULT_MASKED = True
+MASK_TOKEN = "🔒 MASKED (masking policy ON)"
+_policy_cache: dict[str, bool] | None = None
+
 # Datasets we surface in the PII overlay. node_id matches dataflow.py node ids.
 # om_fqn = OpenMetadata table FQN (service.database.schema.table) for tag lookup.
 DATASETS = [
@@ -42,6 +53,44 @@ DATASETS = [
      "om_fqn": "atc_trino.iceberg.curated_masked.sales_orders_masked"},
 ]
 
+KEY_BY_NODE = {ds["node_id"]: ds["key"] for ds in DATASETS}
+DATASET_BY_KEY = {ds["key"]: ds for ds in DATASETS}
+
+
+def _load_policy() -> dict[str, bool]:
+    global _policy_cache
+    if _policy_cache is None:
+        try:
+            _policy_cache = {k: bool(v) for k, v in json.loads(POLICY_PATH.read_text()).items()}
+        except Exception:
+            _policy_cache = {}
+    return _policy_cache
+
+
+def _save_policy(p: dict[str, bool]) -> None:
+    global _policy_cache
+    _policy_cache = p
+    try:
+        POLICY_PATH.parent.mkdir(parents=True, exist_ok=True)
+        POLICY_PATH.write_text(json.dumps(p, indent=2))
+    except Exception:
+        pass
+
+
+def _resolve_key(key_or_node: str) -> str | None:
+    if key_or_node in DATASET_BY_KEY:
+        return key_or_node
+    return KEY_BY_NODE.get(key_or_node)
+
+
+def is_masked(key: str, column: str, *, masked_layer: bool = False) -> bool:
+    """Resolve whether a column is masked: physical masked layer is always masked;
+    otherwise the user policy override wins, falling back to DEFAULT_MASKED."""
+    if masked_layer:
+        return True
+    return _load_policy().get(f"{key}.{column}", DEFAULT_MASKED)
+
+
 # name fragment -> PII category
 PII_RULES: list[tuple[str, str]] = [
     (r"email|e_mail", "EMAIL"),
@@ -66,11 +115,6 @@ def _classify(col: str) -> str | None:
     return None
 
 
-def _is_masked(col: str, dataset_masked_layer: bool) -> bool:
-    c = col.lower()
-    return dataset_masked_layer or c.endswith("_masked") or c.endswith("_hash") or c.endswith("_token")
-
-
 # OpenMetadata PII tag -> our category. OM applies PII.Sensitive / PII.NonSensitive
 # plus optional General/PersonalData tags via auto-classification.
 _OM_CAT = {
@@ -150,9 +194,10 @@ def _build() -> dict[str, Any]:
                 cat = heur or _OM_CAT.get(pii_tag, "PII")
             else:
                 cat = heur
-            masked = _is_masked(c, ds.get("masked_layer", False))
+            masked = is_masked(ds["key"], c, masked_layer=ds.get("masked_layer", False))
             pii_cols.append({
                 "name": c, "category": cat, "masked": masked,
+                "policy_locked": ds.get("masked_layer", False),
                 "source": "openmetadata" if pii_tag else "heuristic",
                 "om_tag": pii_tag,
             })
@@ -184,6 +229,110 @@ def get_pii(use_cache: bool = True) -> dict[str, Any]:
     return data
 
 
+def _trino_query(sql: str, timeout: float = 20.0) -> tuple[list[str], list[list[Any]]]:
+    cols: list[str] = []
+    rows: list[list[Any]] = []
+    with httpx.Client(timeout=timeout) as client:
+        d = client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(), headers={"X-Trino-User": TRINO_USER}).json()
+        while True:
+            if d.get("error"):
+                raise RuntimeError(d["error"].get("message", "trino error"))
+            c = d.get("columns")
+            if c and not cols:
+                cols = [x["name"] for x in c]
+            rows += d.get("data") or []
+            nxt = d.get("nextUri")
+            if not nxt:
+                break
+            d = client.get(nxt).json()
+    return cols, rows
+
+
+class MaskPolicyRequest(BaseModel):
+    key: str            # dataset key or node id
+    column: str         # column name, or "*" for every PII column of the dataset
+    masked: bool
+
+
+class LookupRequest(BaseModel):
+    key: str            # dataset key or node id
+    search: str | None = None
+    limit: int = 5
+
+
 @router.get("")
 async def pii_overview(refresh: bool = False) -> JSONResponse:
     return JSONResponse(get_pii(use_cache=not refresh))
+
+
+@router.get("/policy")
+async def get_policy() -> JSONResponse:
+    return JSONResponse({"ok": True, "default_masked": DEFAULT_MASKED, "policy": _load_policy()})
+
+
+@router.post("/policy")
+async def set_policy(body: MaskPolicyRequest) -> JSONResponse:
+    key = _resolve_key(body.key)
+    if not key:
+        return JSONResponse({"ok": False, "error": f"unknown dataset {body.key}"}, status_code=400)
+    data = get_pii()
+    dset = next((d for d in data["datasets"] if d["key"] == key), None)
+    if body.column == "*":
+        cols = [c["name"] for c in (dset["pii_columns"] if dset else [])]
+    else:
+        cols = [body.column]
+    if not cols:
+        return JSONResponse({"ok": False, "error": "no columns to update"}, status_code=400)
+    p = dict(_load_policy())
+    for col in cols:
+        p[f"{key}.{col}"] = bool(body.masked)
+    _save_policy(p)
+    _cache["data"] = None  # force rebuild so masked flags reflect the new policy
+    return JSONResponse({"ok": True, "key": key, "columns": cols, "masked": bool(body.masked)})
+
+
+@router.post("/lookup")
+async def lookup(body: LookupRequest) -> JSONResponse:
+    """Look up actual records in a dataset, enforcing the masking policy server-side.
+    Masked column values are replaced with MASK_TOKEN and never leave this process."""
+    key = _resolve_key(body.key)
+    if not key:
+        return JSONResponse({"ok": False, "error": f"unknown dataset {body.key}"}, status_code=400)
+    ds = DATASET_BY_KEY[key]
+    dset = next((d for d in get_pii()["datasets"] if d["key"] == key), None)
+    if not dset or not dset["pii_columns"]:
+        return JSONResponse({"ok": False, "error": "no PII columns known for this dataset"}, status_code=400)
+
+    pii_cols = dset["pii_columns"]
+    masked_map = {c["name"]: c["masked"] for c in pii_cols}
+    name_col = next((c["name"] for c in pii_cols if c["category"] == "NAME"), None)
+    select_cols = list(dict.fromkeys(c["name"] for c in pii_cols))
+    col_sql = ", ".join(f'"{c}"' for c in select_cols)
+
+    where = ""
+    if body.search and name_col:
+        safe = body.search.replace("'", "''")
+        where = f" WHERE lower(cast(\"{name_col}\" AS varchar)) LIKE lower('%{safe}%')"
+    limit = max(1, min(body.limit or 5, 25))
+    sql = f"SELECT {col_sql} FROM {ds['table']}{where} LIMIT {limit}"
+
+    try:
+        cols, rows = _trino_query(sql)
+    except Exception as exc:  # noqa: BLE001
+        return JSONResponse({"ok": False, "error": f"query failed: {exc}"}, status_code=502)
+
+    out_rows = []
+    for row in rows:
+        rec: dict[str, Any] = {}
+        for cname, val in zip(cols, row):
+            rec[cname] = MASK_TOKEN if masked_map.get(cname) else val
+        out_rows.append(rec)
+
+    return JSONResponse({
+        "ok": True, "key": key, "table": ds["table"],
+        "rows": out_rows, "row_count": len(out_rows),
+        "masked_columns": [c for c, m in masked_map.items() if m],
+        "unmasked_columns": [c for c, m in masked_map.items() if not m],
+        "note": ("Values shown as the mask token are withheld by the active masking policy "
+                 "and were never sent to the model."),
+    })

ui/src/components/features/DataFlowView.tsx

@@ -1,7 +1,7 @@
 import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react'
-import { GitBranch, Play, RefreshCw, ShieldAlert, ShieldCheck, Loader2, Bot } from 'lucide-react'
+import { GitBranch, Lock, LockOpen, Play, RefreshCw, ShieldAlert, ShieldCheck, Loader2, Bot } from 'lucide-react'
 import type { DataflowEdge, DataflowGraph, DataflowNode } from '../../types'
-import { fetchDataflow, runDataflowMovement, toggleEtlAgent, fetchAgentOpsStatus } from '../../lib/api'
+import { fetchDataflow, runDataflowMovement, toggleEtlAgent, fetchAgentOpsStatus, setPiiMask } from '../../lib/api'
 import { Badge } from '../ui/Badge'
 import { cn } from '../../lib/utils'
 
@@ -158,6 +158,17 @@ export function DataFlowView() {
     setTimeout(() => setRefreshing(false), 400)
   }, [load])
 
+  const [maskBusy, setMaskBusy] = useState<string | null>(null)
+  const onToggleMask = useCallback(async (key: string, column: string, masked: boolean) => {
+    setMaskBusy(`${key}.${column}`)
+    try {
+      await setPiiMask(key, column, masked)
+      await load(true)
+    } finally {
+      setMaskBusy(null)
+    }
+  }, [load])
+
   const piiSummary = graph?.pii_summary ?? {}
   const movementEdges = useMemo(
     () => edges.filter((e) => e.movement_id),
@@ -329,22 +340,57 @@ export function DataFlowView() {
             )}
             {selNode.pii?.has_pii ? (
               <div className="mt-1.5 border-t border-border pt-1.5">
-                <div className="mb-1 flex items-center gap-1 text-[9px] font-medium text-rose-300">
+                <div className="mb-1 flex items-center justify-between gap-1 text-[9px] font-medium text-rose-300">
+                  <span className="flex items-center gap-1">
                     {selNode.pii.all_masked ? <ShieldCheck className="h-3 w-3 text-emerald-300" /> : <ShieldAlert className="h-3 w-3" />}
                     {selNode.pii.pii_count} PII column{selNode.pii.pii_count === 1 ? '' : 's'}
+                  </span>
+                  {!selNode.pii.masked_layer && (
+                    <span className="flex gap-1">
+                      <button
+                        type="button"
+                        onClick={() => onToggleMask(selNode.pii!.key, '*', true)}
+                        className="rounded border border-emerald-400/40 px-1 text-[7px] text-emerald-300 hover:bg-emerald-500/15"
+                      >mask all</button>
+                      <button
+                        type="button"
+                        onClick={() => onToggleMask(selNode.pii!.key, '*', false)}
+                        className="rounded border border-rose-400/40 px-1 text-[7px] text-rose-300 hover:bg-rose-500/15"
+                      >unmask all</button>
+                    </span>
+                  )}
                 </div>
+                <p className="mb-1 text-[7px] leading-tight text-foreground-faint">
+                  Masked columns are hidden from the assistant — it refuses to reveal them.
+                </p>
                 <div className="flex flex-col gap-0.5">
-                  {selNode.pii.columns.map((c) => (
+                  {selNode.pii.columns.map((c) => {
-                    <div key={c.name} className="flex items-center justify-between gap-1 text-[8px]">
+                    const locked = c.policy_locked
+                    const busy = maskBusy === `${selNode.pii!.key}.${c.name}`
+                    return (
+                      <button
+                        key={c.name}
+                        type="button"
+                        disabled={locked || busy}
+                        onClick={() => onToggleMask(selNode.pii!.key, c.name, !c.masked)}
+                        title={locked ? 'Physically masked in the curated layer' : c.masked ? 'Masked — click to allow the assistant to read it' : 'Visible to the assistant — click to mask'}
+                        className={cn(
+                          'flex items-center justify-between gap-1 rounded px-1 py-0.5 text-[8px] transition-colors',
+                          locked ? 'cursor-default opacity-70' : 'hover:bg-surface-overlay',
+                        )}
+                      >
                         <span className="truncate font-mono text-foreground-muted">{c.name}</span>
                         <span className="flex items-center gap-1">
                           <span className="text-blue-300">{c.category}</span>
-                        <span className={c.masked ? 'text-emerald-300' : 'text-rose-300'}>
+                          <span className={cn('inline-flex items-center gap-0.5', c.masked ? 'text-emerald-300' : 'text-rose-300')}>
-                          {c.masked ? 'masked' : 'raw'}
+                            {busy ? <Loader2 className="h-2.5 w-2.5 animate-spin" />
+                              : c.masked ? <Lock className="h-2.5 w-2.5" /> : <LockOpen className="h-2.5 w-2.5" />}
+                            {c.masked ? 'masked' : 'visible'}
                           </span>
                         </span>
-                    </div>
+                      </button>
-                  ))}
+                    )
+                  })}
                 </div>
               </div>
             ) : (

ui/src/lib/api.ts

@@ -155,6 +155,14 @@ export async function fetchPii(refresh = false): Promise<{ datasets: PiiDataset[
   )
 }
 
+export function setPiiMask(key: string, column: string, masked: boolean) {
+  return fetch('/api/pii/policy', {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ key, column, masked }),
+  })
+}
+
 export function toggleEtlAgent(enabled?: boolean) {
   return fetch('/api/agent-ops/etl/toggle', {
     method: 'POST',

ui/src/types.ts

@@ -70,7 +70,7 @@ export type CdcStats = {
   consumed: number
 }
 
-export type PiiColumn = { name: string; category: string; masked: boolean }
+export type PiiColumn = { name: string; category: string; masked: boolean; policy_locked?: boolean }
 
 export type DataflowNode = {
   id: string
@@ -83,6 +83,7 @@ export type DataflowNode = {
   metric: string | null
   url?: string
   pii?: {
+    key: string
     has_pii: boolean
     pii_count: number
     all_masked: boolean