310 lines
14 KiB
Python
310 lines
14 KiB
Python
"""Request-scoped masking manifest and response validation invariants."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
os.environ["KANSHO_DB_PATH"] = str(Path(tempfile.gettempdir()) / "kansho-privacy-manifest-test.sqlite")
|
|
os.environ["KANSHO_FAKE_PROVIDER"] = "1"
|
|
os.environ["KANSHO_FAKE_DETECT"] = "1"
|
|
Path(os.environ["KANSHO_DB_PATH"]).unlink(missing_ok=True)
|
|
|
|
from fastapi.testclient import TestClient
|
|
from identity_store import remember_mapping
|
|
from main import app
|
|
from privacy_gateway import (
|
|
ERROR_EGRESS_VALIDATION,
|
|
IDENTITY_LEAK_RETRY,
|
|
ActiveReplacement,
|
|
GatewayRequest,
|
|
MaskingManifest,
|
|
PrivacyGatewayError,
|
|
_demask,
|
|
_validate_response,
|
|
complete,
|
|
last_trace,
|
|
mask_prompt,
|
|
reset_debug,
|
|
validate_pre_egress,
|
|
)
|
|
from providers import ChatResult
|
|
|
|
|
|
def expect(ok: bool, message: str) -> None:
|
|
if not ok:
|
|
raise SystemExit(f"FAIL: {message}")
|
|
print(f"OK {message}")
|
|
|
|
|
|
def _no_label_keys(payload) -> bool:
|
|
if isinstance(payload, dict):
|
|
if "local_label" in payload:
|
|
return False
|
|
return all(_no_label_keys(value) for value in payload.values())
|
|
if isinstance(payload, list):
|
|
return all(_no_label_keys(item) for item in payload)
|
|
return True
|
|
|
|
|
|
def _run_gateway(profile_id: str, rendered: str, *, purpose: str = "dialogue_turn"):
|
|
return complete(
|
|
GatewayRequest(
|
|
prompt_id="manifest-test",
|
|
purpose=purpose,
|
|
data_class="B",
|
|
profile_id=profile_id,
|
|
payload={"rendered": rendered, "source_text": rendered},
|
|
)
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
reset_debug()
|
|
anna = [{"local_label": "Anna", "token": "PERSON:01", "entity_type": "PERSON"}]
|
|
lars = [{"local_label": "Lars", "token": "PERSON:02", "entity_type": "PERSON"}]
|
|
sushi = [{"local_label": "Sushi", "token": "PERSON:01", "entity_type": "PERSON"}]
|
|
hamburg = [{"local_label": "Hamburg", "token": "PLACE:01", "entity_type": "PLACE"}]
|
|
|
|
inactive = mask_prompt("Heute nur Markt und Kirschen.", anna, "journal_generate")
|
|
expect(inactive.replacements == (), "historical mapping is inactive when absent from the prompt")
|
|
kept = _validate_response("Anna kam später zum Markt.", inactive)
|
|
expect("Anna kam später" in kept, "inactive mapping does not treat the same string as a leak")
|
|
|
|
mixed = mask_prompt("Anna war in Hamburg.", anna + lars + hamburg, "journal_generate")
|
|
active_tokens = {item.token for item in mixed.replacements}
|
|
expect(active_tokens == {"PERSON:01", "PLACE:01"}, "only actually masked mappings are active")
|
|
expect("PERSON:02" not in active_tokens, "unused profile mapping stays inactive")
|
|
|
|
food_only = mask_prompt("Ich esse gerne Sushi.", sushi, "journal_generate")
|
|
expect(food_only.replacements == (), "intentional homonym is not an active identity")
|
|
food_reply = _validate_response("Heute nur von Sushi essen erzählt.", food_only)
|
|
expect("Sushi essen" in food_reply, "unmasked homonym does not block a normal reply")
|
|
|
|
kinship = mask_prompt("Meine Frau Sushi kam später.", sushi, "journal_generate")
|
|
expect(len(kinship.replacements) == 1, "identity mention of the same word is active")
|
|
blocked = False
|
|
try:
|
|
_validate_response("Sushi kam vorbei.", kinship)
|
|
except PrivacyGatewayError as exc:
|
|
blocked = exc.code == "response_validation_failed"
|
|
expect(exc.diagnostics.get("leak_tokens") == ["PERSON:01"], "block reports the token, not the label")
|
|
expect("Sushi" not in json.dumps(exc.diagnostics), "validation diagnostics omit the clear name")
|
|
expect(blocked, "active clear name in the raw reply is a leak")
|
|
homonym_ok = _validate_response("Danach Sushi essen.", kinship)
|
|
expect("Sushi essen" in homonym_ok, "homonym in the reply uses the same classification rule")
|
|
|
|
placeholder = _validate_response("[[PERSON:01]] war ruhig.", kinship)
|
|
expect(placeholder.startswith("[[PERSON:01]]"), "active placeholder passes validation")
|
|
expect(_demask(placeholder, kinship) == "Sushi war ruhig.", "demask happens only after validation")
|
|
leftover = _demask("[[PERSON:01]] und [[PERSON:99]]", kinship)
|
|
expect("Sushi" in leftover and "[[PERSON:99]]" in leftover, "inactive placeholder is not rematerialized")
|
|
expect("Clarissa" not in leftover, "inactive mapping label is not introduced by demask")
|
|
expect("[[" not in IDENTITY_LEAK_RETRY, "retry instruction must not teach bracket placeholders")
|
|
expect(
|
|
_demask("[[ person:01 ]] war ruhig.", kinship) == "Sushi war ruhig.",
|
|
"demask ignores case and inner spacing",
|
|
)
|
|
expect(
|
|
_demask("[[PERSON:1]] war ruhig.", kinship) == "Sushi war ruhig.",
|
|
"demask accepts an unpadded numeric token",
|
|
)
|
|
expect(
|
|
_demask("[[…]] und [[...]] blieben.", kinship) == "[[…]] und [[...]] blieben.",
|
|
"generic ellipsis placeholders are not rematerialized",
|
|
)
|
|
|
|
overlap = mask_prompt(
|
|
"Anna-Lena kam vorbei.",
|
|
[
|
|
{"local_label": "Anna", "token": "PERSON:01"},
|
|
{"local_label": "Anna-Lena", "token": "PERSON:08"},
|
|
],
|
|
"journal_generate",
|
|
)
|
|
expect([item.token for item in overlap.replacements] == ["PERSON:08"], "only the actually replaced overlapping label is active")
|
|
expect("[[PERSON:08]]" in overlap.masked_text, "longer label is the one replaced")
|
|
expect("Anna-Lena" not in overlap.masked_text, "replaced longer label is gone")
|
|
|
|
subject = mask_prompt(
|
|
"Clarissa kaufte Kirschen.",
|
|
[{"local_label": "Clarissa", "token": "PERSON:99"}],
|
|
"journal_generate",
|
|
)
|
|
expect([item.token for item in subject.replacements] == ["PERSON:99"], "person as subject of kaufte remains identity")
|
|
|
|
leaky = MaskingManifest(
|
|
masked_text="Anna war am Markt.",
|
|
available_mapping_count=1,
|
|
replacements=(ActiveReplacement("PERSON:01", "PERSON", 1, "Anna"),),
|
|
)
|
|
pre = False
|
|
try:
|
|
validate_pre_egress("Anna war am Markt.", leaky)
|
|
except PrivacyGatewayError as exc:
|
|
pre = exc.code == ERROR_EGRESS_VALIDATION
|
|
expect(exc.diagnostics.get("pre_egress_validation") == "failed", "pre-egress failure is named")
|
|
expect(pre, "remaining active identity occurrence fails closed before the provider")
|
|
validate_pre_egress("[[PERSON:01]] war am Markt.", leaky)
|
|
|
|
def check(prompt: str, mappings: list[dict], reply: str):
|
|
manifest = mask_prompt(prompt, mappings, "journal_generate")
|
|
try:
|
|
return "ok", _validate_response(reply, manifest), [item.token for item in manifest.replacements]
|
|
except PrivacyGatewayError as exc:
|
|
return exc.code, None, [item.token for item in manifest.replacements]
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
first = pool.submit(check, "Anna kam vorbei.", anna, "Anna kam.")
|
|
second = pool.submit(check, "Heute nur Markt.", anna + lars, "Anna kam.")
|
|
code_a, _text_a, tokens_a = first.result()
|
|
code_b, text_b, tokens_b = second.result()
|
|
expect(code_a == "response_validation_failed", "parallel request A still blocks its own active leak")
|
|
expect(tokens_a == ["PERSON:01"], "parallel request A keeps its own active set")
|
|
expect(code_b == "ok" and "Anna kam" in (text_b or ""), "parallel request B is not blocked by A's mapping")
|
|
expect(tokens_b == [], "parallel request B does not inherit A's active set")
|
|
|
|
with TestClient(app) as client:
|
|
setup = client.post(
|
|
"/api/auth/setup",
|
|
json={"email": "lars@example.test", "name": "Lars", "password": "test-pass"},
|
|
)
|
|
profile_id = setup.json()["profile_id"]
|
|
remember_mapping(profile_id, "Anna", "PERSON:01")
|
|
remember_mapping(profile_id, "Clarissa", "PERSON:99")
|
|
|
|
result = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
|
|
expect(result.allowed, "masked identity request is allowed")
|
|
expect(result.diagnostics.get("available_mapping_count", 0) >= 1, "diagnostics count request mappings")
|
|
expect(result.diagnostics.get("active_mapping_count") >= 1, "diagnostics count actually masked mappings")
|
|
expect("PERSON:01" in (result.diagnostics.get("active_tokens") or []), "diagnostics name active tokens")
|
|
expect(result.trace.get("pre_egress_validation") == "ok", "pre-egress is recorded")
|
|
expect(result.trace.get("response_validation") == "ok", "response validation is recorded")
|
|
expect(_no_label_keys(result.diagnostics), "compact diagnostics have no local_label keys")
|
|
expect(_no_label_keys(result.trace.get("log")), "run log has no local_label keys")
|
|
expect(_no_label_keys(last_trace()), "last compact trace has no local_label keys")
|
|
compact_blob = json.dumps(result.diagnostics)
|
|
expect("local_label" not in compact_blob, "compact diagnostics omit the mapping field name")
|
|
expect("Anna" not in compact_blob, "compact diagnostics omit the clear name")
|
|
|
|
provider_calls = {"n": 0}
|
|
|
|
def boom(_messages, _policy):
|
|
provider_calls["n"] += 1
|
|
raise AssertionError("provider must not be called after pre-egress failure")
|
|
|
|
def leaky_mask(rendered, mappings, purpose):
|
|
return MaskingManifest(
|
|
masked_text=rendered,
|
|
available_mapping_count=len(mappings or []),
|
|
replacements=(ActiveReplacement("PERSON:01", "PERSON", 1, "Anna"),),
|
|
)
|
|
|
|
blocked_egress = False
|
|
with patch("privacy_gateway.mask_prompt", leaky_mask), patch("privacy_gateway.complete_model", boom):
|
|
try:
|
|
_run_gateway(profile_id, "Anna war am Markt.")
|
|
except PrivacyGatewayError as exc:
|
|
blocked_egress = exc.code == ERROR_EGRESS_VALIDATION
|
|
expect(blocked_egress, "pre-egress failure aborts the request")
|
|
expect(provider_calls["n"] == 0, "provider is not called when an active identity remains")
|
|
|
|
seen = []
|
|
|
|
def leak_then_ok(messages, _policy):
|
|
seen.append(messages[0].get("content") or "")
|
|
if len(seen) == 1:
|
|
return ChatResult(
|
|
content="Anna stand den ganzen Nachmittag am Markt.",
|
|
model="fake",
|
|
usage={"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18, "cost": 0.01},
|
|
context_compression="not_applicable",
|
|
)
|
|
return ChatResult(
|
|
content="[[PERSON:01]] stand den ganzen Nachmittag am Markt.",
|
|
model="fake",
|
|
usage={"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20, "cost": 0.02},
|
|
context_compression="not_applicable",
|
|
)
|
|
|
|
with patch("privacy_gateway.complete_model", leak_then_ok):
|
|
repaired = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
|
|
expect(len(seen) == 2, "a real leak retries exactly once")
|
|
expect(IDENTITY_LEAK_RETRY in seen[1], "retry uses the generic correction instruction")
|
|
expect("Anna stand den ganzen Nachmittag" not in seen[1], "discarded raw reply is not part of the retry prompt")
|
|
expect("[[PERSON:01]]" in seen[0] and "Anna" not in seen[0].replace("[[PERSON:01]]", ""), "retry keeps the already masked prompt")
|
|
expect("Anna" in (repaired.content or ""), "successful retry is demasked locally")
|
|
expect(repaired.diagnostics.get("response_validation_retry") == 1, "retry is recorded without the raw reply")
|
|
|
|
def spaced_token(_messages, _policy):
|
|
return ChatResult(
|
|
content="[[ person:01 ]] stand den ganzen Nachmittag am Markt.",
|
|
model="fake",
|
|
usage={},
|
|
context_compression="not_applicable",
|
|
)
|
|
|
|
with patch("privacy_gateway.complete_model", spaced_token):
|
|
restored = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
|
|
expect(
|
|
"Anna stand den ganzen Nachmittag" in (restored.content or ""),
|
|
"gateway demasks case and spacing variants",
|
|
)
|
|
expect("[[" not in (restored.content or ""), "no leftover placeholder after a successful demask")
|
|
expect(repaired.diagnostics.get("prompt_tokens") == 22, "retry aggregates prompt tokens")
|
|
expect(repaired.diagnostics.get("completion_tokens") == 16, "retry aggregates completion tokens")
|
|
expect(repaired.diagnostics.get("total_tokens") == 38, "retry aggregates total tokens")
|
|
expect(abs(float(repaired.diagnostics.get("cost") or 0) - 0.03) < 1e-9, "retry aggregates cost")
|
|
|
|
inactive_calls = {"n": 0}
|
|
|
|
def keep_inactive(messages, _policy):
|
|
inactive_calls["n"] += 1
|
|
return ChatResult(
|
|
content="Clarissa blieb heute unerwähnt am Hafen.",
|
|
model="fake",
|
|
usage={},
|
|
context_compression="not_applicable",
|
|
)
|
|
|
|
with patch("privacy_gateway.complete_model", keep_inactive):
|
|
kept_inactive = _run_gateway(profile_id, "Heute nur Markt und Kirschen am Hafen.")
|
|
expect(inactive_calls["n"] == 1, "inactive mapping does not start a retry")
|
|
expect(kept_inactive.trace.get("response_validation_retry") in (None, 0), "inactive mapping does not mark a retry")
|
|
expect("Clarissa blieb heute unerwähnt" in (kept_inactive.content or ""), "inactive mapping keeps the model text")
|
|
|
|
retry_fail = []
|
|
|
|
def always_leak(messages, _policy):
|
|
retry_fail.append(messages[0].get("content") or "")
|
|
return ChatResult(
|
|
content="Anna stand den ganzen Nachmittag am Markt.",
|
|
model="fake",
|
|
usage={},
|
|
context_compression="not_applicable",
|
|
)
|
|
|
|
failed = False
|
|
with patch("privacy_gateway.complete_model", always_leak):
|
|
try:
|
|
_run_gateway(profile_id, "Ich war mit Anna am Markt.")
|
|
except PrivacyGatewayError as exc:
|
|
failed = exc.code == "response_validation_failed"
|
|
expect(exc.diagnostics.get("response_validation_retry") == 1, "failed retry is still only one extra call")
|
|
expect("Anna stand" not in json.dumps(exc.diagnostics.get("log") or []), "discarded reply is not persisted in the log")
|
|
expect(failed, "second leak stays fail-closed")
|
|
expect(len(retry_fail) == 2, "gateway stops after one retry")
|
|
|
|
print("All privacy manifest tests passed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|