fix(pii): keep context scoring for thresholded long texts and match data: case-insensitively

- Long-text NER windows also run the pattern entities when a score_threshold is set, so context-boosted matches clear it as they do on short texts
- The data: URI scheme is matched case-insensitively
- e2e asserts each name is masked as one <PERSON> span and is registered as test:pii-redaction:e2e

Author: waleedlatif1

apps/pii/server.py

@@ -243,11 +243,18 @@ def _exceeds_nlp_length(text: str, language: str) -> bool:
     return limit is not None and len(text) > limit
 
 
-def _ner_window_entities(entities: list[str] | None, language: str) -> list[str]:
-    """The requested entities only spaCy NER produces and this language serves, or
-    all of them when `entities` is empty (detect-all)."""
-    ner = NER_ENTITIES & set(analyzer.get_supported_entities(language))
-    return sorted(ner if not entities else ner.intersection(entities))
+def _window_entities(
+    entities: list[str] | None, language: str, score_threshold: float | None
+) -> list[str]:
+    """The requested built-in entities this language serves that windows must find,
+    or all of them when `entities` is empty (detect-all). Without a score_threshold
+    that is only what spaCy NER produces; with one it also includes the pattern
+    entities, whose context-boosted score can decide whether a match clears the
+    threshold and needs the real NLP artifacts only a window has."""
+    served = set(analyzer.get_supported_entities(language))
+    if score_threshold is None or score_threshold <= 0:
+        served &= NER_ENTITIES
+    return sorted(served if not entities else served.intersection(entities))
 
 
 def _analyze_long(
@@ -263,13 +270,14 @@ def _analyze_long(
     Pattern, checksum, and custom recognizers have no length limit, so they run
     once over the whole text (against blank NLP artifacts, as the regex-only path
     does) and no match can straddle a cut. Only NER runs in overlapping windows
-    (see NER_WINDOW_CHARS), asked for just its own entities. A span touching a
+    (see NER_WINDOW_CHARS), asked for just its own entities, plus the pattern
+    entities when a score_threshold makes context scoring matter (see
+    `_window_entities`). A span touching a
     window's far edge may be cut short, so it is left to the next window, which
     sees it whole; a span ending inside the overlap before that edge was already
     found by the previous window and is skipped. The rest are shifted to
-    whole-text offsets, and overlapping spans are left to the anonymizer, which
-    merges them. Context words cannot raise a pattern match's score here, which
-    matters only when a score_threshold is set."""
+    whole-text offsets; a span found by both passes keeps its higher score, and
+    overlapping spans are left to the anonymizer, which merges them."""
     entities = entities or None
     results = analyzer.analyze(
         text=text,
@@ -280,7 +288,7 @@ def _analyze_long(
         nlp_artifacts=_BLANK_ARTIFACTS.get(language),
         ad_hoc_recognizers=ad_hoc_recognizers or None,
     )
-    window_entities = _ner_window_entities(entities, language)
+    window_entities = _window_entities(entities, language, score_threshold)
     if not window_entities:
         return results
     step = NER_WINDOW_CHARS - NER_WINDOW_OVERLAP

apps/sim/lib/logs/execution/pii-redaction.test.ts

@@ -170,6 +170,7 @@ describe('encoded binary content', () => {
     ['wrapped at 76 columns', wrapped(76, '\n')],
     ['wrapped at 64 columns with CRLF', `${wrapped(64, '\r\n')}\r\n`],
     ['as a wrapped data URL', `data:application/pdf;base64,${wrapped(76, '\n')}`],
+    ['as a data URL with an uppercase scheme', `DATA:application/pdf;base64,${encodedFile}`],
     [
       'as a data URL with long metadata',
       `data:application/pdf;name=${'quarterly-report-'.repeat(40)}.pdf;base64,${encodedFile}`,

apps/sim/lib/logs/execution/pii-redaction.ts

@@ -60,7 +60,7 @@ function base64Body(value: string): string | null {
  * masking cannot check such a value for PII, which lives inside the decoded file.
  */
 function isEncodedBinary(value: string): boolean {
-  if (value.startsWith('data:')) {
+  if (/^data:/i.test(value)) {
     const comma = value.indexOf(',')
     return (
       comma !== -1 &&

apps/sim/package.json

@@ -25,6 +25,7 @@
     "test:scim:e2e": "bun run scripts/test-scim-e2e.ts",
     "test:sim-links:e2e": "bun run scripts/test-sim-links-e2e.ts",
     "test:people:e2e": "bun run scripts/test-people-e2e.ts",
+    "test:pii-redaction:e2e": "bun --no-env-file scripts/test-pii-redaction-e2e.ts",
     "test:browser-extension:e2e": "bun --no-env-file scripts/test-browser-extension-e2e.ts",
     "test:isolated-vm-remote:e2e": "bun run scripts/test-isolated-vm-remote-e2e.ts",
     "test:mobile:e2e": "bun --no-env-file scripts/test-mobile-e2e.ts",

apps/sim/scripts/test-pii-redaction-e2e.ts

@@ -91,9 +91,10 @@ function assertNameMasked(masked: unknown, failedMarker: string): void {
   assert.equal(typeof masked, 'string', 'Expected the long text back as a string')
   const text = masked as string
   assert(!text.includes(failedMarker), 'The long text was scrubbed instead of masked')
-  assert(!text.includes('Ada Lovelace'), 'The name at the end of the long text was not masked')
-  assert(text.endsWith('before Friday.'), 'The end of the long text was lost')
-  assert(text.includes('<PERSON>'), 'Expected a <PERSON> token where the name was')
+  assert(
+    text.endsWith('Please send questions to <PERSON> before Friday.'),
+    'The name at the end of the long text was not masked as one <PERSON> span'
+  )
 }
 
 async function main(): Promise<void> {
@@ -166,8 +167,11 @@ async function main(): Promise<void> {
       `Presidio answered ${response.status}: ${body.slice(0, 200)}`
     )
     const [masked] = (JSON.parse(body) as { texts: string[] }).texts
-    assert(!masked.includes('Grace Hopper'), 'The name across the first NER window edge survived')
-    assert(!masked.includes('Ada Lovelace'), 'The name at the end survived')
+    assert(
+      masked.includes(`${lead}<PERSON> before Friday.`),
+      'The name across the first NER window edge was not masked as one <PERSON> span'
+    )
+    assert(masked.endsWith('to <PERSON> before Friday.'), 'The name at the end survived')
     assert(!masked.includes('EMP 123456'), 'The custom pattern spanning whitespace survived')
     assert(masked.includes('<EMPLOYEE_ID>'), 'Expected the custom replacement')
     return { textChars: text.length }