fix(pii): mask offloaded large payloads chunk-by-chunk instead of aborting at 16MB

apps/pii/server.py

@@ -176,23 +176,50 @@ def build_analyzer() -> AnalyzerEngine:
 batch_analyzer = BatchAnalyzerEngine(analyzer_engine=analyzer)
 anonymizer = AnonymizerEngine()
 
-# Every entity the spaCy NER recognizers can produce. A request touching any of
-# these must run spaCy; a request naming only non-NER (regex/checksum) entities can
-# skip it. Derived from the live registry so it stays authoritative if Presidio's
-# default entity set changes (e.g. ORGANIZATION), unioned with a known floor so an
-# unexpectedly empty derivation can never let an NER request skip the NLP pass.
+# Every entity the spaCy NER recognizers can actually produce. A request touching
+# any of these must run spaCy; a request naming only non-NER (regex/checksum)
+# entities can skip it. The registry's SpacyRecognizers CLAIM every entity in
+# Presidio's default NER-model mapping — including PHONE_NUMBER/AGE/ID/EMAIL,
+# which exist for transformer models and which no spaCy model can emit — so the
+# claimed set is intersected with the entities the loaded models' NER labels map
+# onto. Without that filter, selecting PHONE_NUMBER (present in nearly every
+# redaction rule) silently forces the full spaCy pass and the regex-only fast
+# path never fires. The floor guarantees the core NER entities always take the
+# full pass even if label introspection ever derives empty.
 _SPACY_NER_FLOOR = frozenset({"PERSON", "LOCATION", "NRP", "DATE_TIME", "ORGANIZATION"})
-NER_ENTITIES = _SPACY_NER_FLOOR | frozenset(
-    entity
-    for recognizer in analyzer.registry.recognizers
-    if isinstance(recognizer, SpacyRecognizer)
-    for entity in recognizer.supported_entities
+
+
+def _producible_ner_entities() -> frozenset:
+    """Presidio entities some loaded spaCy model's NER labels map onto."""
+    configuration = getattr(analyzer.nlp_engine, "ner_model_configuration", None)
+    mapping = configuration.model_to_presidio_entity_mapping if configuration else {}
+    producible = set()
+    for nlp in getattr(analyzer.nlp_engine, "nlp", {}).values():
+        if "ner" not in nlp.pipe_names:
+            continue
+        for label in nlp.get_pipe("ner").labels:
+            entity = mapping.get(label)
+            if entity:
+                producible.add(entity)
+    return frozenset(producible)
+
+
+NER_ENTITIES = _SPACY_NER_FLOOR | (
+    frozenset(
+        entity
+        for recognizer in analyzer.registry.recognizers
+        if isinstance(recognizer, SpacyRecognizer)
+        for entity in recognizer.supported_entities
+    )
+    & _producible_ner_entities()
 )
 
 # One blank NlpArtifacts per language, built once at startup. Passing these to
 # analyze() skips nlp_engine.process_text (the spaCy tok2vec+ner pass) entirely:
-# the pattern recognizers still match on the raw text, SpacyRecognizer is excluded
-# by the entity filter, and score_threshold is unset so detection is identical.
+# the pattern recognizers still match on the raw text, SpacyRecognizer finds
+# nothing in blank artifacts (and can only be consulted at all for claimed-but-
+# unproducible entities like PHONE_NUMBER), and score_threshold is unset so
+# detection is identical.
 # Only context-based score boosting (which needs real tokens) is unavailable — an
 # accepted trade for skipping NER on the hot block-output path. Read-only, so it
 # is safe to share across requests and workers.

apps/sim/lib/guardrails/mask-client.test.ts

@@ -3,13 +3,15 @@
  */
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockToken, mockBaseUrl } = vi.hoisted(() => ({
+const { mockToken, mockBaseUrl, mockSleep } = vi.hoisted(() => ({
   mockToken: vi.fn(),
   mockBaseUrl: vi.fn(),
+  mockSleep: vi.fn(),
 }))
 
 vi.mock('@/lib/auth/internal', () => ({ generateInternalToken: mockToken }))
 vi.mock('@/lib/core/utils/urls', () => ({ getInternalApiBaseUrl: mockBaseUrl }))
+vi.mock('@sim/utils/helpers', () => ({ sleep: mockSleep }))
 
 import { maskPIIBatchViaHttp } from '@/lib/guardrails/mask-client'
 
@@ -20,6 +22,7 @@ describe('maskPIIBatchViaHttp', () => {
     vi.clearAllMocks()
     mockToken.mockResolvedValue('tok')
     mockBaseUrl.mockReturnValue('http://app.internal:3000')
+    mockSleep.mockResolvedValue(undefined)
     fetchMock = vi.fn(async (_url: string, init: { body: string }) => {
       const { texts } = JSON.parse(init.body) as { texts: string[] }
       return new Response(JSON.stringify({ masked: texts.map((t) => `M(${t})`) }), {
@@ -52,10 +55,80 @@ describe('maskPIIBatchViaHttp', () => {
     expect(fetchMock).toHaveBeenCalledTimes(3) // 2000-per-request cap
   })
 
-  it('throws on a non-2xx response so the caller can scrub', async () => {
-    fetchMock.mockResolvedValueOnce(new Response('boom', { status: 500 }))
+  it('throws immediately on a deterministic 4xx without retrying', async () => {
+    fetchMock.mockResolvedValueOnce(new Response('bad request', { status: 400 }))
 
     await expect(maskPIIBatchViaHttp(['a'], [])).rejects.toThrow(/mask-batch request failed/)
+    expect(fetchMock).toHaveBeenCalledTimes(1)
+    expect(mockSleep).not.toHaveBeenCalled()
+  })
+
+  it('retries a transient 5xx with backoff and then succeeds', async () => {
+    fetchMock.mockResolvedValueOnce(new Response('deploying', { status: 503 }))
+
+    const out = await maskPIIBatchViaHttp(['a'], [])
+
+    expect(out).toEqual(['M(a)'])
+    expect(fetchMock).toHaveBeenCalledTimes(2)
+    expect(mockSleep).toHaveBeenCalledTimes(1)
+  })
+
+  it('retries a rejected fetch (network error) and then succeeds', async () => {
+    fetchMock.mockRejectedValueOnce(new TypeError('fetch failed'))
+
+    const out = await maskPIIBatchViaHttp(['a'], [])
+
+    expect(out).toEqual(['M(a)'])
+    expect(fetchMock).toHaveBeenCalledTimes(2)
+  })
+
+  it('retries a runtime-level request timeout and then succeeds', async () => {
+    const timeout = new Error('The operation timed out.')
+    timeout.name = 'TimeoutError'
+    fetchMock.mockRejectedValueOnce(timeout)
+
+    const out = await maskPIIBatchViaHttp(['a'], [])
+
+    expect(out).toEqual(['M(a)'])
+    expect(fetchMock).toHaveBeenCalledTimes(2)
+  })
+
+  it('gives up after the retry budget is exhausted on a persistent 5xx', async () => {
+    fetchMock.mockImplementation(async () => new Response('down', { status: 503 }))
+
+    await expect(maskPIIBatchViaHttp(['a'], [])).rejects.toThrow(/mask-batch request failed/)
+    expect(fetchMock).toHaveBeenCalledTimes(8)
+    expect(mockSleep).toHaveBeenCalledTimes(7)
+  })
+
+  it('mints a fresh internal token per attempt', async () => {
+    fetchMock.mockResolvedValueOnce(new Response('deploying', { status: 503 }))
+
+    await maskPIIBatchViaHttp(['a'], [])
+
+    expect(mockToken).toHaveBeenCalledTimes(2)
+  })
+
+  it('does not retry a null 200 body (deterministic, not a transient TypeError)', async () => {
+    fetchMock.mockResolvedValueOnce(
+      new Response('null', { status: 200, headers: { 'content-type': 'application/json' } })
+    )
+
+    await expect(maskPIIBatchViaHttp(['a'], [])).rejects.toThrow(/unexpected result/)
+    expect(fetchMock).toHaveBeenCalledTimes(1)
+    expect(mockSleep).not.toHaveBeenCalled()
+  })
+
+  it('does not retry a shape mismatch (deterministic server bug)', async () => {
+    fetchMock.mockResolvedValueOnce(
+      new Response(JSON.stringify({ nope: true }), {
+        status: 200,
+        headers: { 'content-type': 'application/json' },
+      })
+    )
+
+    await expect(maskPIIBatchViaHttp(['a'], [])).rejects.toThrow(/unexpected result/)
+    expect(fetchMock).toHaveBeenCalledTimes(1)
   })
 
   it('returns [] without any request for empty input', async () => {

apps/sim/lib/guardrails/mask-client.ts

@@ -1,3 +1,5 @@
+import { sleep } from '@sim/utils/helpers'
+import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
 import type { GuardrailsMaskBatchResult } from '@/lib/api/contracts'
 import { generateInternalToken } from '@/lib/auth/internal'
 import { env } from '@/lib/core/config/env'
@@ -18,6 +20,60 @@ import type { CustomPiiPattern } from '@/lib/guardrails/pii-entities'
  */
 const CHUNK_CONCURRENCY = env.PII_MASK_CHUNK_CONCURRENCY ?? 64
 
+/**
+ * Per-chunk retry budget for transient failures (network errors, 408/429/5xx).
+ * A large payload fans out into many chunk requests, so a single blip — an ALB
+ * 502 during a deploy, a Presidio pod restart — must not fail the whole
+ * redaction (and, on the execution-altering stages, abort the run). With the
+ * default 500ms→30s jittered backoff this rides out ~2 minutes of outage per
+ * chunk before giving up. Deterministic failures (4xx, shape mismatches) throw
+ * immediately.
+ */
+const MAX_CHUNK_ATTEMPTS = 8
+
+const RETRYABLE_STATUSES = new Set([408, 429, 500, 502, 503, 504])
+
+class MaskChunkHttpError extends Error {
+  constructor(
+    message: string,
+    readonly status: number,
+    readonly retryAfterMs: number | null
+  ) {
+    super(message)
+    this.name = 'MaskChunkHttpError'
+  }
+}
+
+function isRetryableChunkError(error: unknown): boolean {
+  if (error instanceof MaskChunkHttpError) {
+    return RETRYABLE_STATUSES.has(error.status)
+  }
+  // A rejected fetch (connection refused/reset, DNS, socket drop) is transient —
+  // Node wraps these in TypeError('fetch failed'). Runtime-level request
+  // timeouts (undici's default 300s headers/body timeout, Bun's TimeoutError)
+  // and mid-flight socket closes surface with their own names/codes per runtime;
+  // all are congestion or connection churn, not a deterministic failure: a chunk
+  // queued behind a saturated Presidio fleet must retry, not fail the payload.
+  if (error instanceof TypeError) {
+    return true
+  }
+  const { name, code } = (error ?? {}) as { name?: unknown; code?: unknown }
+  if (name === 'TimeoutError' || name === 'HeadersTimeoutError' || name === 'BodyTimeoutError') {
+    return true
+  }
+  return (
+    typeof code === 'string' &&
+    [
+      'ECONNRESET',
+      'ECONNREFUSED',
+      'EPIPE',
+      'ETIMEDOUT',
+      'ConnectionClosed',
+      'ConnectionRefused',
+    ].includes(code)
+  )
+}
+
 /**
  * Mask PII across many strings via the internal app-container endpoint.
  *
@@ -29,8 +85,10 @@ const CHUNK_CONCURRENCY = env.PII_MASK_CHUNK_CONCURRENCY ?? 64
  * concurrency, so a large payload fans out rather than serializing; order is
  * preserved, so the returned array matches `texts` length.
  *
- * Rejects on any non-2xx, timeout, or shape mismatch so the caller can apply
- * its own fail-safe (scrubbing rather than leaking).
+ * Transient chunk failures (network errors, 408/429/5xx) retry with jittered
+ * backoff (see {@link MAX_CHUNK_ATTEMPTS}); only a deterministic failure or an
+ * exhausted retry budget rejects, so the caller can apply its own fail-safe
+ * (scrubbing rather than leaking).
  */
 export async function maskPIIBatchViaHttp(
   texts: string[],
@@ -64,8 +122,29 @@ async function postChunk(
   language: string | undefined,
   customPatterns: CustomPiiPattern[] | undefined
 ): Promise<string[]> {
-  // Mint per request: a single token (5min TTL) can expire mid-batch when a
-  // large execution fans out into many sequential chunk requests.
+  for (let attempt = 1; ; attempt++) {
+    try {
+      return await postChunkOnce(url, texts, entityTypes, language, customPatterns)
+    } catch (error) {
+      if (attempt >= MAX_CHUNK_ATTEMPTS || !isRetryableChunkError(error)) {
+        throw error
+      }
+      const retryAfterMs = error instanceof MaskChunkHttpError ? error.retryAfterMs : null
+      await sleep(backoffWithJitter(attempt, retryAfterMs))
+    }
+  }
+}
+
+async function postChunkOnce(
+  url: string,
+  texts: string[],
+  entityTypes: string[],
+  language: string | undefined,
+  customPatterns: CustomPiiPattern[] | undefined
+): Promise<string[]> {
+  // Mint per attempt: a single token (5min TTL) can expire mid-batch when a
+  // large execution fans out into many sequential chunk requests or a chunk
+  // spends its retry budget waiting out an outage.
   const token = await generateInternalToken()
 
   // boundary-raw-fetch: internal server-to-server call to the app container (internal JWT auth, configurable base URL)
@@ -80,11 +159,15 @@ async function postChunk(
 
   if (!response.ok) {
     const detail = await response.text().catch(() => '')
-    throw new Error(`PII mask-batch request failed (${response.status}): ${detail.slice(0, 200)}`)
+    throw new MaskChunkHttpError(
+      `PII mask-batch request failed (${response.status}): ${detail.slice(0, 200)}`,
+      response.status,
+      parseRetryAfter(response.headers.get('retry-after'))
+    )
   }
 
-  const data = (await response.json()) as GuardrailsMaskBatchResult
-  if (!Array.isArray(data.masked)) {
+  const data = (await response.json()) as GuardrailsMaskBatchResult | null
+  if (!data || !Array.isArray(data.masked)) {
     throw new Error('PII mask-batch returned an unexpected result')
   }
   return data.masked