fix(pii): mask offloaded large payloads chunk-by-chunk and retry transient mask failures

A block output past the 16MB inline materialization ceiling aborted the run
before masking even started: the redaction path hydrated the whole offloaded
value at once, and the pre-flight size assert fired on the manifest's total
byteSize. Large-array manifests now page one stored chunk at a time
(materialize -> mask -> re-store, rebuilt via the manifest writer with preview
derived from masked items), so peak heap stays ~one chunk regardless of payload
size. Single refs up to the 64MB durable cap hydrate with a raised budget and
run serially outside the concurrency pool.

Mask-batch chunk requests now retry transient failures (network errors,
408/429/5xx, honoring Retry-After) with jittered backoff, so a single ALB blip
or Presidio pod restart no longer fails a whole payload's redaction. Nested-ref
masking now runs the string pass before ref substitution, fixing a latent
double-mask when a masked nested value shrinks back inline.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01A1JYstmLHk9qMGyBDqYRcJ

Author: TheodoreSpeaks

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,59 @@ 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('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 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,39 @@ 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;
+  // anything else (shape mismatch, token minting failure) is deterministic.
+  return error instanceof TypeError
+}
+
 /**
  * Mask PII across many strings via the internal app-container endpoint.
  *
@@ -29,8 +64,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 +101,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,7 +138,11 @@ 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