@@ -8,7 +8,7 @@ import { generateInternalToken } from '@/lib/auth/internal'
88import { env } from '@/lib/core/config/env'
99import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
1010import { getInternalApiBaseUrl } from '@/lib/core/utils/urls'
11- import { chunkIndicesByBudget } from '@/lib/guardrails/pii-batching'
11+ import { chunkIndicesByBudget , joinSegments , segmentTexts } from '@/lib/guardrails/pii-batching'
1212import type { CustomPiiPattern } from '@/lib/guardrails/pii-entities'
1313
1414const logger = createLogger ( 'PiiMaskClient' )
@@ -36,6 +36,15 @@ const CHUNK_CONCURRENCY = env.PII_MASK_CHUNK_CONCURRENCY ?? 64
3636 */
3737const MAX_CHUNK_ATTEMPTS = 8
3838
39+ /**
40+ * Attempts for a chunk the route answers with a 500. The route keeps 503 for a
41+ * Presidio outage and 422 for input Presidio refused, so a 500 is an unclassified
42+ * failure that may be deterministic for this input: two retries cover a one-off
43+ * crash without spending the full outage budget resending input that fails the
44+ * same way every time.
45+ */
46+ const MAX_ATTEMPTS_AFTER_SERVER_ERROR = 3
47+
3948const RETRYABLE_STATUSES = new Set ( [ 408 , 429 , 500 , 502 , 503 , 504 ] )
4049
4150class MaskChunkHttpError extends Error {
@@ -97,9 +106,9 @@ function isOutageError(error: unknown): boolean {
97106
98107interface MaskPIIBatchOptions {
99108 /**
100- * When set, a chunk that cannot be masked has each of its strings replaced with
101- * this value instead of rejecting the call, so one bad chunk costs only its own
102- * strings. After an outage (see {@link OUTAGE_STATUSES}) exhausts a chunk's
109+ * When set, a chunk that cannot be masked has every string with a piece in it
110+ * replaced with this value instead of rejecting the call, so one bad chunk costs
111+ * only its own strings. After an outage (see {@link OUTAGE_STATUSES}) exhausts a chunk's
103112 * retries, chunks not yet sent are replaced the same way.
104113 */
105114 failedChunkPlaceholder ?: string
@@ -118,7 +127,9 @@ interface MaskPIIBatchResult {
118127 * Only the app task reaches the Presidio service (it holds `PII_URL`), but the
119128 * log-redaction persist path also runs inside the trigger.dev runtime — so
120129 * redaction always routes through HTTP, the same way the guardrails tool does.
121- * Strings are grouped into byte/count-budgeted chunks (keeping each request far
130+ * A string too long for one NER pass is cut into pieces first (see
131+ * {@link segmentTexts}), so no request carries a text Presidio would refuse.
132+ * Pieces are grouped into byte/count-budgeted chunks (keeping each request far
122133 * under the 10MB Next body limit) and the chunks are sent with bounded
123134 * concurrency, so a large payload fans out rather than serializing; order is
124135 * preserved, so the returned `masked` array matches `texts` length.
@@ -141,38 +152,47 @@ export async function maskPIIBatchViaHttp(
141152
142153 const url = `${ getInternalApiBaseUrl ( ) } /api/guardrails/mask-batch`
143154 const placeholder = options . failedChunkPlaceholder
144- const masked = new Array < string > ( texts . length )
155+ const { segments, owners } = segmentTexts ( texts )
156+ const maskedSegments = new Array < string > ( segments . length )
157+ const scrubbedTexts = new Set < number > ( )
145158 let serviceUnavailable = false
146- let scrubbedCount = 0
147159 let firstFailure : unknown
148160
149- const scrubChunk = ( indices : number [ ] , value : string ) => {
150- for ( const i of indices ) masked [ i ] = value
151- scrubbedCount += indices . length
161+ const scrubChunk = ( indices : number [ ] ) => {
162+ for ( const k of indices ) scrubbedTexts . add ( owners [ k ] )
152163 }
153164
154- await mapWithConcurrency ( chunkIndicesByBudget ( texts ) , CHUNK_CONCURRENCY , async ( indices ) => {
165+ await mapWithConcurrency ( chunkIndicesByBudget ( segments ) , CHUNK_CONCURRENCY , async ( indices ) => {
155166 if ( serviceUnavailable && placeholder !== undefined ) {
156- scrubChunk ( indices , placeholder )
167+ scrubChunk ( indices )
157168 return
158169 }
159170 try {
160- const chunk = indices . map ( ( i ) => texts [ i ] )
171+ const chunk = indices . map ( ( k ) => segments [ k ] )
161172 const out = await postChunk ( url , chunk , entityTypes , language , customPatterns )
162173 if ( out . length !== chunk . length ) {
163174 throw new Error ( 'PII mask-batch returned an unexpected result' )
164175 }
165- indices . forEach ( ( originalIndex , k ) => {
166- masked [ originalIndex ] = out [ k ]
176+ indices . forEach ( ( segmentIndex , k ) => {
177+ maskedSegments [ segmentIndex ] = out [ k ]
167178 } )
168179 } catch ( error ) {
169180 if ( placeholder === undefined ) throw error
170181 if ( isOutageError ( error ) ) serviceUnavailable = true
171182 firstFailure ??= error
172- scrubChunk ( indices , placeholder )
183+ scrubChunk ( indices )
173184 }
174185 } )
175186
187+ const joined = joinSegments ( maskedSegments , owners , texts . length )
188+ // A text with any unmasked piece is replaced whole: its other pieces are masked,
189+ // but keeping them would persist a fragment around the hole.
190+ const masked =
191+ placeholder === undefined
192+ ? joined
193+ : joined . map ( ( text , i ) => ( scrubbedTexts . has ( i ) ? placeholder : text ) )
194+ const scrubbedCount = scrubbedTexts . size
195+
176196 if ( scrubbedCount > 0 ) {
177197 logger . error ( 'PII masking failed for some chunks; replaced their strings' , {
178198 error : getErrorMessage ( firstFailure ) ,
@@ -196,7 +216,11 @@ async function postChunk(
196216 try {
197217 return await postChunkOnce ( url , texts , entityTypes , language , customPatterns )
198218 } catch ( error ) {
199- if ( attempt >= MAX_CHUNK_ATTEMPTS || ! isRetryableChunkError ( error ) ) {
219+ const maxAttempts =
220+ error instanceof MaskChunkHttpError && error . status === 500
221+ ? MAX_ATTEMPTS_AFTER_SERVER_ERROR
222+ : MAX_CHUNK_ATTEMPTS
223+ if ( attempt >= maxAttempts || ! isRetryableChunkError ( error ) ) {
200224 throw error
201225 }
202226 const retryAfterMs = error instanceof MaskChunkHttpError ? error . retryAfterMs : null
0 commit comments