"""Merge the deterministic and fuzzy layers into the final entity list.
Rules:
- * deterministic spans (score 1.0) always outrank fuzzy spans on overlap;
+ * deterministic spans (score 1.0) always outrank fuzzy spans on overlap,
+ while uncovered fuzzy fragments remain eligible;
* among fuzzy spans, the longer one wins, then the higher score;
* fuzzy spans below `score_threshold` are dropped (deterministic = 1.0 always
passes);
from .entities import Span, overlaps
+def _subtract_overlaps(span: Span, blockers: list[Span]) -> list[Span]:
+ """Return the fragments of `span` not covered by any blocker."""
+ intervals = [(span.start, span.end)]
+
+ for blocker in blockers:
+ uncovered: list[tuple[int, int]] = []
+ for start, end in intervals:
+ if blocker.end <= start or end <= blocker.start:
+ uncovered.append((start, end))
+ continue
+ if start < blocker.start:
+ uncovered.append((start, blocker.start))
+ if blocker.end < end:
+ uncovered.append((blocker.end, end))
+ intervals = uncovered
+
+ return [Span(span.entity_type, start, end, span.score) for start, end in intervals]
+
+
def merge(
deterministic: list[Span],
fuzzy: list[Span],
# Deterministic spans are pre-resolved (non-overlapping) and take precedence.
accepted: list[Span] = [s for s in deterministic if s.score >= score_threshold]
+ deterministic_blockers = list(accepted)
+ uncovered_fuzzy = [
+ fragment for span in fuzzy for fragment in _subtract_overlaps(span, deterministic_blockers)
+ ]
+
# Longer fuzzy spans first, then higher score, for stable overlap resolution.
- for span in sorted(fuzzy, key=lambda s: (-s.length, -s.score)):
+ for span in sorted(uncovered_fuzzy, key=lambda s: (-s.length, -s.score)):
if span.score < score_threshold:
continue
if any(overlaps(span, a) for a in accepted):
from detector.merge import merge
-def test_deterministic_precedence_over_overlapping_fuzzy():
- det = [Span(IBAN_CODE, 0, 27, 1.0)]
- fuzzy = [Span(LOCATION, 0, 4, 0.9)] # e.g. "DE89" mis-tagged
+def test_deterministic_full_overlap_removes_fuzzy():
+ det = [Span(IBAN_CODE, 10, 20, 1.0)]
+ fuzzy = [Span(PERSON, 12, 18, 0.9)]
out = merge(det, fuzzy, None, 0.7)
- assert out == [Span(IBAN_CODE, 0, 27, 1.0)]
+ assert out == [Span(IBAN_CODE, 10, 20, 1.0)]
+
+
+def test_deterministic_left_overlap_preserves_fuzzy_left_fragment():
+ det = [Span(IBAN_CODE, 10, 20, 1.0)]
+ fuzzy = [Span(PERSON, 0, 15, 0.9)]
+ out = merge(det, fuzzy, None, 0.7)
+ assert out == [Span(PERSON, 0, 10, 0.9), Span(IBAN_CODE, 10, 20, 1.0)]
+
+
+def test_deterministic_right_overlap_preserves_fuzzy_right_fragment():
+ det = [Span(IBAN_CODE, 10, 20, 1.0)]
+ fuzzy = [Span(PERSON, 15, 30, 0.9)]
+ out = merge(det, fuzzy, None, 0.7)
+ assert out == [Span(IBAN_CODE, 10, 20, 1.0), Span(PERSON, 20, 30, 0.9)]
+
+
+def test_deterministic_split_overlap_preserves_both_fuzzy_fragments():
+ det = [Span(IBAN_CODE, 10, 20, 1.0)]
+ fuzzy = [Span(PERSON, 0, 30, 0.9)]
+ out = merge(det, fuzzy, None, 0.7)
+ assert out == [
+ Span(PERSON, 0, 10, 0.9),
+ Span(IBAN_CODE, 10, 20, 1.0),
+ Span(PERSON, 20, 30, 0.9),
+ ]
def test_fuzzy_below_threshold_dropped():
findDenylistedEntities,
mergeDenylistEntities,
PIIDetector,
+ subtractPlaceholderOverlaps,
} from "./detect";
const originalFetch = globalThis.fetch;
});
});
+ describe("subtractPlaceholderOverlaps", () => {
+ const text = "left [[TOKEN_1]] right";
+ const placeholders = ["[[TOKEN_1]]"];
+
+ test("removes an entity fully inside a placeholder", () => {
+ const entities = [{ entity_type: "PERSON", start: 7, end: 14, score: 0.91 }];
+
+ expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([]);
+ });
+
+ test("leaves non-overlapping entities unchanged", () => {
+ const entities = [
+ { entity_type: "PERSON", start: 0, end: 4, score: 0.91 },
+ { entity_type: "LOCATION", start: 17, end: 22, score: 0.82 },
+ ];
+
+ expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual(entities);
+ });
+
+ test("keeps the uncovered left fragment of a partial overlap", () => {
+ const entities = [{ entity_type: "PERSON", start: 0, end: 10, score: 0.91 }];
+
+ expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([
+ { entity_type: "PERSON", start: 0, end: 5, score: 0.91 },
+ ]);
+ });
+
+ test("keeps the uncovered right fragment of a partial overlap", () => {
+ const entities = [{ entity_type: "PERSON", start: 10, end: 22, score: 0.91 }];
+
+ expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([
+ { entity_type: "PERSON", start: 16, end: 22, score: 0.91 },
+ ]);
+ });
+
+ test("splits an entity around a placeholder", () => {
+ const entities = [{ entity_type: "PERSON", start: 0, end: 22, score: 0.91 }];
+
+ expect(subtractPlaceholderOverlaps(text, entities, placeholders)).toEqual([
+ { entity_type: "PERSON", start: 0, end: 5, score: 0.91 },
+ { entity_type: "PERSON", start: 16, end: 22, score: 0.91 },
+ ]);
+ });
+
+ test("subtracts multiple placeholders from one spanning entity", () => {
+ const multiText = "aa [[ONE_1]] bb [[TWO_2]] cc";
+ const entities = [{ entity_type: "LOCATION", start: 1, end: 27, score: 0.83 }];
+
+ expect(subtractPlaceholderOverlaps(multiText, entities, ["[[ONE_1]]", "[[TWO_2]]"])).toEqual([
+ { entity_type: "LOCATION", start: 1, end: 3, score: 0.83 },
+ { entity_type: "LOCATION", start: 12, end: 16, score: 0.83 },
+ { entity_type: "LOCATION", start: 25, end: 27, score: 0.83 },
+ ]);
+ });
+ });
+
describe("filterAllowlistedEntities", () => {
test("filters entities matching allowlist pattern", () => {
const text = "You are Claude Code, Anthropic's official CLI for Claude.";
index = text.indexOf(placeholder, index + placeholder.length);
}
}
- return spans;
+ return spans.sort((a, b) => a.start - b.start || a.end - b.end);
+}
+
+/**
+ * Remove placeholder-covered intervals from detector entities.
+ *
+ * Entity offsets stay relative to the original text. A partial overlap can
+ * therefore produce both a left and a right fragment with the original type
+ * and score.
+ */
+export function subtractPlaceholderOverlaps(
+ text: string,
+ entities: PIIEntity[],
+ knownPlaceholders: readonly string[],
+): PIIEntity[] {
+ if (entities.length === 0 || knownPlaceholders.length === 0) return entities;
+
+ const placeholders = placeholderSpans(text, knownPlaceholders);
+ if (placeholders.length === 0) return entities;
+
+ return entities.flatMap((entity) => {
+ let fragments = [entity];
+
+ for (const placeholder of placeholders) {
+ fragments = fragments.flatMap((fragment) => {
+ if (!overlaps(fragment, placeholder)) return [fragment];
+
+ const uncovered: PIIEntity[] = [];
+ if (fragment.start < placeholder.start) {
+ uncovered.push({ ...fragment, end: placeholder.start });
+ }
+ if (placeholder.end < fragment.end) {
+ uncovered.push({ ...fragment, start: placeholder.end });
+ }
+ return uncovered;
+ });
+
+ if (fragments.length === 0) break;
+ }
+
+ return fragments;
+ });
}
export function findDenylistedEntities(
const matches = denylist.flatMap(({ pattern, type, regex }) =>
regex ? findRegexMatches(text, pattern, type) : findLiteralMatches(text, pattern, type),
);
- if (matches.length === 0 || knownPlaceholders.length === 0) return matches;
-
- // Drop matches inside an already-masked placeholder; re-masking its internals would corrupt the earlier mask.
- const masked = placeholderSpans(text, knownPlaceholders);
- if (masked.length === 0) return matches;
- return matches.filter((m) => !masked.some((p) => overlaps(m, p)));
+ return subtractPlaceholderOverlaps(text, matches, knownPlaceholders);
}
// Additive merge: a denylist match extends coverage but never shrinks an overlapping detector span; overlaps become their union and the detector type wins.
const denylistedEntities = findDenylistedEntities(span.text, denylist, knownPlaceholders);
const detectedEntities = config.pii_detection.enabled ? await this.detectPII(span.text) : [];
- const filteredEntities = filterAllowlistedEntities(span.text, detectedEntities, allowlist);
+ const uncoveredEntities = subtractPlaceholderOverlaps(
+ span.text,
+ detectedEntities,
+ knownPlaceholders,
+ );
+ const filteredEntities = filterAllowlistedEntities(span.text, uncoveredEntities, allowlist);
spanEntities.push(mergeDenylistEntities(filteredEntities, denylistedEntities));
}
expect(knownPlaceholders).toEqual(["[[API_KEY_SK_1]]"]);
});
+ test("masks sensitive text beside a secret placeholder without remasking the placeholder", async () => {
+ const connectionString = "postgres://admin:S3cretPass@db.example.com:5432/appdb";
+ const detector = new PIIDetector();
+ detector.detectPII = mock(async (text: string) => {
+ expect(text).toBe("hunter2 [[CONNECTION_STRING_1]]");
+ return [{ entity_type: "PERSON", start: 0, end: text.length, score: 0.93 }];
+ });
+ mockAnalyzeRequest.mockImplementationOnce((maskedRequest, extractor, knownPlaceholders) =>
+ detector.analyzeRequest(
+ maskedRequest as OpenAIRequest,
+ extractor as typeof openaiExtractor,
+ knownPlaceholders,
+ ),
+ );
+
+ const result = await processPrivacyPipeline(
+ request(`hunter2 ${connectionString}`),
+ {
+ ...baseConfig,
+ secrets_detection: {
+ ...baseConfig.secrets_detection,
+ entities: ["CONNECTION_STRING"],
+ },
+ },
+ openaiExtractor,
+ );
+
+ expect(result.requestAfterSecrets.messages[0].content).toBe("hunter2 [[CONNECTION_STRING_1]]");
+ expect(result.piiResult?.detection.spanEntities[0]).toEqual([
+ { entity_type: "PERSON", start: 0, end: 8, score: 0.93 },
+ ]);
+ expect(result.request.messages[0].content).toBe("[[PERSON_1]][[CONNECTION_STRING_1]]");
+ expect(result.piiMaskingContext?.mapping["[[PERSON_1]]"]).toBe("hunter2 ");
+ expect(result.secretsResult.maskingContext?.mapping["[[CONNECTION_STRING_1]]"]).toBe(
+ connectionString,
+ );
+ });
+
test("returns privacy facts without route decisions", async () => {
const result = await processPrivacyPipeline(request("Hello"), baseConfig, openaiExtractor);
}
});
+ test("masks a detector fragment beside a secret placeholder without corrupting it", async () => {
+ const connectionString = "postgres://admin:S3cretPass@db.example.com:5432/appdb";
+ mockDetectPII.mockImplementationOnce((text: string) => {
+ expect(text).toBe("hunter2 [[CONNECTION_STRING_1]]");
+ return Promise.resolve([
+ {
+ entity_type: "PERSON",
+ start: 0,
+ end: text.length,
+ score: 0.93,
+ },
+ ]);
+ });
+
+ const res = await app.request("/api/mask", {
+ method: "POST",
+ headers: { "Content-Type": "application/json" },
+ body: JSON.stringify({ text: `hunter2 ${connectionString}` }),
+ });
+
+ expect(res.status).toBe(200);
+ const body = (await res.json()) as {
+ masked: string;
+ context: Record<string, string>;
+ entities: { type: string; placeholder: string }[];
+ };
+ expect(body.masked).toBe("[[PERSON_1]][[CONNECTION_STRING_1]]");
+ expect(body.context["[[PERSON_1]]"]).toBe("hunter2 ");
+ expect(body.context["[[CONNECTION_STRING_1]]"]).toBe(connectionString);
+ expect(body.entities).toContainEqual({
+ type: "CONNECTION_STRING",
+ placeholder: "[[CONNECTION_STRING_1]]",
+ });
+ expect(body.entities).toContainEqual({ type: "PERSON", placeholder: "[[PERSON_1]]" });
+ });
+
test("returns 400 for malformed JSON", async () => {
const res = await app.request("/api/mask", {
method: "POST",
findDenylistedEntities,
getPIIDetector,
mergeDenylistEntities,
+ subtractPlaceholderOverlaps,
} from "../pii/detect";
import { mask as maskPII } from "../pii/mask";
import { detectSecrets } from "../secrets/detect";
const piiEntities = config.pii_detection.enabled ? await detector.detectPII(maskedText) : [];
scanTimeMs = Date.now() - piiStartTime;
- const filteredEntities = filterAllowlistedEntities(
+ const uncoveredEntities = subtractPlaceholderOverlaps(
maskedText,
piiEntities,
+ Object.keys(context.mapping),
+ );
+ const filteredEntities = filterAllowlistedEntities(
+ maskedText,
+ uncoveredEntities,
config.masking.allowlist,
);
const entitiesToMask = mergeDenylistEntities(
git.99rst.org / sgasser-llm-shield.git / commitdiff