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