diff --git a/apps/pii/server.py b/apps/pii/server.py index fdd7c1b2604..05b23e9b0b1 100644 --- a/apps/pii/server.py +++ b/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. diff --git a/apps/sim/lib/guardrails/mask-client.test.ts b/apps/sim/lib/guardrails/mask-client.test.ts index 3fcb1894353..f8ee031545c 100644 --- a/apps/sim/lib/guardrails/mask-client.test.ts +++ b/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 () => { diff --git a/apps/sim/lib/guardrails/mask-client.ts b/apps/sim/lib/guardrails/mask-client.ts index 7f5776b8524..7972e0e55cc 100644 --- a/apps/sim/lib/guardrails/mask-client.ts +++ b/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 { - // 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 { + // 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 diff --git a/apps/sim/lib/logs/execution/pii-large-values.test.ts b/apps/sim/lib/logs/execution/pii-large-values.test.ts index 15292ef1419..3176b1131e3 100644 --- a/apps/sim/lib/logs/execution/pii-large-values.test.ts +++ b/apps/sim/lib/logs/execution/pii-large-values.test.ts @@ -1,33 +1,41 @@ /** * @vitest-environment node */ +import { sleep } from '@sim/utils/helpers' import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockMaterializeRef, mockCompact, mockMaterializeManifest, mockMaskBatch } = vi.hoisted( - () => ({ - mockMaterializeRef: vi.fn(), - mockCompact: vi.fn(), - mockMaterializeManifest: vi.fn(), - mockMaskBatch: vi.fn(), - }) -) +const { mockMaterializeRef, mockStoreLargeValue, mockCompact, mockMaskBatch } = vi.hoisted(() => ({ + mockMaterializeRef: vi.fn(), + mockStoreLargeValue: vi.fn(), + mockCompact: vi.fn(), + mockMaskBatch: vi.fn(), +})) vi.mock('@/lib/execution/payloads/store', () => ({ materializeLargeValueRef: mockMaterializeRef, + storeLargeValue: mockStoreLargeValue, })) vi.mock('@/lib/execution/payloads/serializer', () => ({ compactExecutionPayload: mockCompact, })) -vi.mock('@/lib/execution/payloads/large-array-manifest', () => ({ - materializeLargeArrayManifest: mockMaterializeManifest, -})) -vi.mock('@/lib/execution/payloads/large-array-manifest-metadata', () => ({ - isLargeArrayManifest: () => false, -})) +vi.mock('@/lib/execution/payloads/materialization.server', () => { + const MAX_INLINE_MATERIALIZATION_BYTES = 16 * 1024 * 1024 + return { + MAX_INLINE_MATERIALIZATION_BYTES, + MAX_DURABLE_LARGE_VALUE_BYTES: 64 * 1024 * 1024, + assertInlineMaterializationSize: (size: number, maxBytes?: number) => { + if (size > (maxBytes ?? MAX_INLINE_MATERIALIZATION_BYTES)) { + throw new Error('Execution memory limit exceeded. Reduce payload size and try again.') + } + }, + } +}) vi.mock('@/lib/guardrails/mask-client', () => ({ maskPIIBatchViaHttp: mockMaskBatch, })) +import type { LargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest' +import { isLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest-metadata' import { redactLargeValueRefs, redactLargeValueRefsInValue, @@ -44,16 +52,58 @@ const REF = { const STORE = { workspaceId: 'ws-1', workflowId: 'wf-1', executionId: 'ex-1', userId: 'u-1' } -describe('redactLargeValueRefs', () => { - beforeEach(() => { - vi.clearAllMocks() - mockMaskBatch.mockImplementation(async (texts: string[]) => texts.map((t) => `MASKED(${t})`)) - // compact echoes its input so we can assert the masked content is what's re-stored. - mockCompact.mockImplementation(async (value: unknown) => value) +/** Chunk arrays keyed by ref id, served by mockMaterializeRef and captured by mockStoreLargeValue. */ +const chunkData = new Map() +let storedRefCounter = 0 + +function makeChunkRef(id: string, size: number, kind = 'array') { + return { __simLargeValueRef: true, version: 1, id, kind, size } as const +} + +function makeManifest( + chunks: Array<{ id: string; size: number; items: unknown[] }> +): LargeArrayManifest { + for (const chunk of chunks) { + chunkData.set(chunk.id, chunk.items) + } + return { + __simLargeArrayManifest: true, + version: 2, + kind: 'array', + totalCount: chunks.reduce((sum, c) => sum + c.items.length, 0), + chunkCount: chunks.length, + byteSize: chunks.reduce((sum, c) => sum + c.size, 0), + chunks: chunks.map((c) => ({ + ref: makeChunkRef(c.id, c.size), + count: c.items.length, + byteSize: c.size, + })), + preview: [{ note: 'raw unmasked preview' }], + } +} + +function installDefaultMocks() { + mockMaskBatch.mockImplementation(async (texts: string[]) => texts.map((t) => `MASKED(${t})`)) + // compact echoes its input so we can assert the masked content is what's re-stored. + mockCompact.mockImplementation(async (value: unknown) => value) + mockMaterializeRef.mockImplementation(async (ref: { id: string }) => chunkData.get(ref.id)) + mockStoreLargeValue.mockImplementation(async (value: unknown, _json: string, size: number) => { + const id = `lv_${String(storedRefCounter++).padStart(12, '0')}` + chunkData.set(id, value) + return { __simLargeValueRef: true, version: 1, id, kind: 'array', size } }) +} + +beforeEach(() => { + vi.clearAllMocks() + chunkData.clear() + storedRefCounter = 0 + installDefaultMocks() +}) +describe('redactLargeValueRefs', () => { it('hydrates, masks, and re-stores a large-value ref (content preserved, PII masked)', async () => { - mockMaterializeRef.mockResolvedValue({ note: 'contact bob', id: 42 }) + chunkData.set(REF.id, { note: 'contact bob', id: 42 }) const result = await redactLargeValueRefs( { finalOutput: REF }, @@ -62,7 +112,11 @@ describe('redactLargeValueRefs', () => { expect(mockMaterializeRef).toHaveBeenCalledWith( REF, - expect.objectContaining({ executionId: 'ex-1', trackReference: false }) + expect.objectContaining({ + executionId: 'ex-1', + trackReference: false, + maxBytes: 64 * 1024 * 1024, + }) ) expect(result.finalOutput).toEqual({ note: 'MASKED(contact bob)', id: 42 }) expect(mockCompact).toHaveBeenCalledTimes(1) @@ -79,7 +133,7 @@ describe('redactLargeValueRefs', () => { }) it('falls back to the marker when re-store throws (never leaks)', async () => { - mockMaterializeRef.mockResolvedValue({ note: 'secret@x.com' }) + chunkData.set(REF.id, { note: 'secret@x.com' }) mockCompact.mockRejectedValueOnce(new Error('s3 down')) const result = await redactLargeValueRefs( { finalOutput: REF }, @@ -91,9 +145,8 @@ describe('redactLargeValueRefs', () => { it('hydrates+masks refs across multiple payload keys (parallel, cross-key)', async () => { const refA = { ...REF, id: 'lv_aaaaaaaaaaaa' } const refB = { ...REF, id: 'lv_bbbbbbbbbbbb' } - mockMaterializeRef.mockImplementation(async (ref: { id: string }) => - ref.id === 'lv_aaaaaaaaaaaa' ? { note: 'call bob' } : { note: 'email amy' } - ) + chunkData.set(refA.id, { note: 'call bob' }) + chunkData.set(refB.id, { note: 'email amy' }) const result = await redactLargeValueRefs( { finalOutput: refA, traceSpans: [{ output: refB }] }, @@ -110,9 +163,7 @@ describe('redactLargeValueRefs', () => { const refA = { ...REF, id: 'lv_aaaaaaaaaaaa' } const refB = { ...REF, id: 'lv_bbbbbbbbbbbb' } // refA materializes fine; refB can't — in throw mode the whole redaction must abort. - mockMaterializeRef.mockImplementation(async (ref: { id: string }) => - ref.id === 'lv_aaaaaaaaaaaa' ? { note: 'ok' } : undefined - ) + chunkData.set(refA.id, { note: 'ok' }) await expect( redactLargeValueRefs( @@ -132,17 +183,246 @@ describe('redactLargeValueRefs', () => { expect(result).toEqual(payload) expect(mockMaterializeRef).not.toHaveBeenCalled() }) + + it('masks a single ref past the inline ceiling with the raised durable budget', async () => { + const bigRef = { ...REF, id: 'lv_bigbigbigbig', size: 30 * 1024 * 1024 } + chunkData.set(bigRef.id, { note: 'ssn 123-45-6789' }) + + const result = await redactLargeValueRefs( + { finalOutput: bigRef }, + { entityTypes: [], language: 'en', store: STORE, onFailure: 'throw' } + ) + + expect(mockMaterializeRef).toHaveBeenCalledWith( + bigRef, + expect.objectContaining({ maxBytes: 64 * 1024 * 1024 }) + ) + expect(result.finalOutput).toEqual({ note: 'MASKED(ssn 123-45-6789)' }) + }) + + it('processes oversized single refs serially, after the pooled refs', async () => { + const smallRef = { ...REF, id: 'lv_smallsmall12' } + const bigRef = { ...REF, id: 'lv_bigbigbigbig', size: 30 * 1024 * 1024 } + chunkData.set(smallRef.id, { note: 'small' }) + chunkData.set(bigRef.id, { note: 'big' }) + + await redactLargeValueRefs( + { finalOutput: { a: bigRef, b: smallRef } }, + { entityTypes: [], language: 'en', store: STORE } + ) + + const order = mockMaterializeRef.mock.calls.map(([ref]) => (ref as { id: string }).id) + expect(order).toEqual(['lv_smallsmall12', 'lv_bigbigbigbig']) + }) + + it('completes an oversized ref whose content nests another oversized ref (no gate deadlock)', async () => { + const outer = { ...REF, id: 'lv_outerbig1234', size: 30 * 1024 * 1024 } + const inner = { ...REF, id: 'lv_innerbig1234', size: 20 * 1024 * 1024 } + chunkData.set(outer.id, { note: 'outer bob', deep: inner }) + chunkData.set(inner.id, { note: 'inner amy' }) + + const result = await redactLargeValueRefs( + { finalOutput: outer }, + { entityTypes: [], language: 'en', store: STORE, onFailure: 'throw' } + ) + + expect(result.finalOutput).toEqual({ + note: 'MASKED(outer bob)', + deep: { note: 'MASKED(inner amy)' }, + }) + }) + + it('serializes nested oversized refs discovered under different pooled parents', async () => { + const parentA = { ...REF, id: 'lv_parentaaaaaa' } + const parentB = { ...REF, id: 'lv_parentbbbbbb' } + const nestedA = { ...REF, id: 'lv_nestedaaaaaa', size: 30 * 1024 * 1024 } + const nestedB = { ...REF, id: 'lv_nestedbbbbbb', size: 30 * 1024 * 1024 } + chunkData.set(parentA.id, { deep: nestedA }) + chunkData.set(parentB.id, { deep: nestedB }) + let inFlight = 0 + let maxInFlight = 0 + mockMaterializeRef.mockImplementation(async (ref: { id: string; size: number }) => { + if (ref.size > 16 * 1024 * 1024) { + inFlight++ + maxInFlight = Math.max(maxInFlight, inFlight) + await sleep(1) + inFlight-- + } + return chunkData.get(ref.id) + }) + chunkData.set(nestedA.id, { note: 'a' }) + chunkData.set(nestedB.id, { note: 'b' }) + + const result = await redactLargeValueRefs( + { finalOutput: { a: parentA, b: parentB } }, + { entityTypes: [], language: 'en', store: STORE } + ) + + expect((result.finalOutput as any).a).toEqual({ deep: { note: 'MASKED(a)' } }) + expect((result.finalOutput as any).b).toEqual({ deep: { note: 'MASKED(b)' } }) + expect(maxInFlight).toBe(1) + }) + + it('processes a manifest containing an oversized chunk serially, after the pooled refs', async () => { + const smallRef = { ...REF, id: 'lv_smallsmall12' } + chunkData.set(smallRef.id, { note: 'small' }) + const bigChunkManifest = makeManifest([ + { id: 'lv_hugechunk000', size: 20 * 1024 * 1024, items: [{ note: 'one giant item' }] }, + ]) + + await redactLargeValueRefs( + { finalOutput: { a: bigChunkManifest, b: smallRef } }, + { entityTypes: [], language: 'en', store: STORE } + ) + + const order = mockMaterializeRef.mock.calls.map(([ref]) => (ref as { id: string }).id) + expect(order).toEqual(['lv_smallsmall12', 'lv_hugechunk000']) + }) }) -describe('redactLargeValueRefsInValue (arbitrary blockStates)', () => { - beforeEach(() => { - vi.clearAllMocks() - mockMaskBatch.mockImplementation(async (texts: string[]) => texts.map((t) => `MASKED(${t})`)) - mockCompact.mockImplementation(async (value: unknown) => value) +describe('redactManifest — chunk-wise', () => { + it('masks a >16MB manifest chunk-by-chunk without tripping the inline ceiling', async () => { + const manifest = makeManifest([ + { id: 'lv_chunk0chunk0', size: 9_000_000, items: [{ note: 'alpha' }, { note: 'beta' }] }, + { id: 'lv_chunk1chunk1', size: 9_000_000, items: [{ note: 'gamma' }, { note: 'delta' }] }, + { id: 'lv_chunk2chunk2', size: 9_000_000, items: [{ note: 'omega' }] }, + ]) + + const result = await redactLargeValueRefs( + { finalOutput: manifest }, + { entityTypes: ['PERSON'], language: 'en', store: STORE, onFailure: 'throw' } + ) + + const masked = result.finalOutput as LargeArrayManifest + expect(isLargeArrayManifest(masked)).toBe(true) + expect(masked.totalCount).toBe(5) + + const maskedItems = masked.chunks.flatMap((chunk) => chunkData.get(chunk.ref.id) as unknown[]) + expect(maskedItems).toEqual([ + { note: 'MASKED(alpha)' }, + { note: 'MASKED(beta)' }, + { note: 'MASKED(gamma)' }, + { note: 'MASKED(delta)' }, + { note: 'MASKED(omega)' }, + ]) + // Manifests re-store through the manifest writer, never the whole-value serializer. + expect(mockCompact).not.toHaveBeenCalled() + // One hydration per source chunk, read-only and with the raised per-chunk budget. + const chunkReads = mockMaterializeRef.mock.calls.filter(([ref]) => + String((ref as { id: string }).id).startsWith('lv_chunk') + ) + expect(chunkReads).toHaveLength(3) + for (const [, context] of chunkReads) { + expect(context).toMatchObject({ trackReference: false, maxBytes: 64 * 1024 * 1024 }) + } + }) + + it('derives the rebuilt preview from masked items, never the raw source preview', async () => { + const manifest = makeManifest([ + { id: 'lv_chunk0chunk0', size: 9_000_000, items: [{ note: 'bob smith' }] }, + { id: 'lv_chunk1chunk1', size: 9_000_000, items: [{ note: 'amy jones' }] }, + ]) + + const result = await redactLargeValueRefs( + { finalOutput: manifest }, + { entityTypes: ['PERSON'], language: 'en', store: STORE } + ) + + const masked = result.finalOutput as LargeArrayManifest + expect(masked.preview).toEqual([{ note: 'MASKED(bob smith)' }]) + }) + + it('returns an empty manifest unchanged in shape for a zero-item manifest', async () => { + const empty: LargeArrayManifest = { + __simLargeArrayManifest: true, + version: 2, + kind: 'array', + totalCount: 0, + chunkCount: 0, + byteSize: 0, + chunks: [], + preview: [], + } + + const result = await redactLargeValueRefs( + { finalOutput: empty }, + { entityTypes: [], language: 'en', store: STORE } + ) + + const masked = result.finalOutput as LargeArrayManifest + expect(isLargeArrayManifest(masked)).toBe(true) + expect(masked.totalCount).toBe(0) + expect(mockStoreLargeValue).not.toHaveBeenCalled() + }) + + it('recursively masks a nested large-value ref inside a chunk item', async () => { + const nestedRef = { ...REF, id: 'lv_nestednested' } + chunkData.set(nestedRef.id, { note: 'nested bob' }) + const manifest = makeManifest([ + { id: 'lv_chunk0chunk0', size: 9_000_000, items: [{ note: 'top', deep: nestedRef }] }, + ]) + + const result = await redactLargeValueRefs( + { finalOutput: manifest }, + { entityTypes: [], language: 'en', store: STORE } + ) + + const masked = result.finalOutput as LargeArrayManifest + const [item] = chunkData.get(masked.chunks[0].ref.id) as Array> + expect(item.note).toBe('MASKED(top)') + expect(item.deep).toEqual({ note: 'MASKED(nested bob)' }) + }) + + describe('partial chunk failure', () => { + it('scrubs the whole manifest to the marker in scrub mode', async () => { + const manifest = makeManifest([ + { id: 'lv_chunk0chunk0', size: 9_000_000, items: [{ note: 'ok' }] }, + { id: 'lv_chunk1chunk1', size: 9_000_000, items: [{ note: 'lost' }] }, + ]) + chunkData.delete('lv_chunk1chunk1') + + const result = await redactLargeValueRefs( + { finalOutput: manifest }, + { entityTypes: [], language: 'en', store: STORE } + ) + + expect(result.finalOutput).toBe('[REDACTION_FAILED]') + }) + + it('throws PiiRedactionError in throw mode', async () => { + const manifest = makeManifest([ + { id: 'lv_chunk0chunk0', size: 9_000_000, items: [{ note: 'ok' }] }, + { id: 'lv_chunk1chunk1', size: 9_000_000, items: [{ note: 'lost' }] }, + ]) + chunkData.delete('lv_chunk1chunk1') + + await expect( + redactLargeValueRefs( + { finalOutput: manifest }, + { entityTypes: [], language: 'en', store: STORE, onFailure: 'throw' } + ) + ).rejects.toBeInstanceOf(PiiRedactionError) + }) + }) + + it('fails fast when a chunk materializes with the wrong item count', async () => { + const manifest = makeManifest([ + { id: 'lv_chunk0chunk0', size: 9_000_000, items: [{ note: 'a' }, { note: 'b' }] }, + ]) + chunkData.set('lv_chunk0chunk0', [{ note: 'a' }]) + + const result = await redactLargeValueRefs( + { finalOutput: manifest }, + { entityTypes: [], language: 'en', store: STORE } + ) + + expect(result.finalOutput).toBe('[REDACTION_FAILED]') }) +}) +describe('redactLargeValueRefsInValue (arbitrary blockStates)', () => { it('hydrates + re-stores a ref nested in a non-RedactablePayload shape', async () => { - mockMaterializeRef.mockResolvedValue({ note: 'contact bob' }) + chunkData.set(REF.id, { note: 'contact bob' }) const blockStates = { 'block-1': { output: REF }, 'block-2': { output: { plain: 'hi' } } } const result = await redactLargeValueRefsInValue(blockStates, { @@ -167,7 +447,7 @@ describe('redactLargeValueRefsInValue (arbitrary blockStates)', () => { }) it('rethrows a re-store failure as PiiRedactionError under throw mode', async () => { - mockMaterializeRef.mockResolvedValue({ note: 'secret@x.com' }) + chunkData.set(REF.id, { note: 'secret@x.com' }) mockCompact.mockRejectedValueOnce(new Error('s3 down')) await expect( diff --git a/apps/sim/lib/logs/execution/pii-large-values.ts b/apps/sim/lib/logs/execution/pii-large-values.ts index 1797964a065..dd323eb6b46 100644 --- a/apps/sim/lib/logs/execution/pii-large-values.ts +++ b/apps/sim/lib/logs/execution/pii-large-values.ts @@ -3,9 +3,17 @@ import { getErrorMessage } from '@sim/utils/errors' import { env } from '@/lib/core/config/env' import { mapWithConcurrency } from '@/lib/core/utils/concurrency' import type { LargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest' -import { materializeLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest' +import { + appendLargeArrayManifest, + createLargeArrayManifest, + readLargeArrayManifestSlice, +} from '@/lib/execution/payloads/large-array-manifest' import { isLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest-metadata' import { isLargeValueRef, type LargeValueRef } from '@/lib/execution/payloads/large-value-ref' +import { + MAX_DURABLE_LARGE_VALUE_BYTES, + MAX_INLINE_MATERIALIZATION_BYTES, +} from '@/lib/execution/payloads/materialization.server' import { compactExecutionPayload } from '@/lib/execution/payloads/serializer' import type { LargeValueStoreContext } from '@/lib/execution/payloads/store' import { materializeLargeValueRef } from '@/lib/execution/payloads/store' @@ -35,6 +43,41 @@ export interface RedactLargeValueRefsOptions { * than feeding a marker (or leaking raw bytes) downstream. */ onFailure?: PiiRedactionFailureMode + /** + * @internal Shared across every nested invocation spawned from one entry-point + * call, so oversized hydrations serialize globally — not just within their own + * {@link resolveReplacements} pass. Created by the entry points; callers never + * set it. + */ + oversizedGate?: OversizedHydrationGate + /** + * @internal True while running under the gate. Nested oversized work inside a + * gated ref runs directly instead of re-acquiring — the holder would otherwise + * wait on itself (deadlock). Safe because the holder is globally serialized. + */ + withinOversizedGate?: boolean +} + +/** Promise-chain mutex for oversized hydrations (see {@link requiresSerialHydration}). */ +interface OversizedHydrationGate { + chain: Promise +} + +/** Queue `fn` behind every previously gated task; failures don't poison the chain. */ +async function runSerially(gate: OversizedHydrationGate, fn: () => Promise): Promise { + const run = gate.chain.then(fn) + gate.chain = run.then( + () => undefined, + () => undefined + ) + return run +} + +/** Ensure one shared gate exists for the whole entry-point call tree. */ +function withOversizedGate(options: RedactLargeValueRefsOptions): RedactLargeValueRefsOptions { + return options.oversizedGate + ? options + : { ...options, oversizedGate: { chain: Promise.resolve() } } } /** @@ -57,6 +100,7 @@ export async function redactLargeValueRefs( payload: RedactablePayload, options: RedactLargeValueRefsOptions ): Promise { + const gatedOptions = withOversizedGate(options) // Collect refs across the WHOLE payload first (shared `seen`), so every ref is // hydrated+masked+re-stored in one bounded-concurrency pass instead of one // sequential pass per key. A ref shared across keys is walked/masked once. @@ -67,7 +111,7 @@ export async function redactLargeValueRefs( } if (refs.length === 0) return payload - const replacements = await resolveReplacements(refs, options) + const replacements = await resolveReplacements(refs, gatedOptions) const result: RedactablePayload = { ...payload } for (const key of Object.keys(payload) as (keyof RedactablePayload)[]) { if (payload[key] !== undefined) result[key] = substituteRefs(payload[key], replacements) @@ -85,7 +129,7 @@ export async function redactLargeValueRefsInValue( value: T, options: RedactLargeValueRefsOptions ): Promise { - return (await redactValueRefs(value, options)) as T + return (await redactValueRefs(value, withOversizedGate(options))) as T } /** Sync-collect every ref/manifest in `value`, then async-replace each, then sync-substitute. */ @@ -108,10 +152,30 @@ async function redactValueRefs( */ const REF_CONCURRENCY = env.PII_REF_CONCURRENCY ?? 4 +/** + * A single ref past the default inline ceiling hydrates its whole blob at once + * (~2-3× its serialized size in transient heap) — and a manifest whose packer + * emitted an oversized chunk (a single item larger than the chunk target) + * hydrates that chunk the same way. Those run serially instead of in the + * {@link REF_CONCURRENCY} pool — one 64MB hydration is fine, four at once is an + * OOM on the small trigger.dev machines. Normally-chunked manifests page ~8MB + * at a time and stay pooled regardless of total size. + */ +function requiresSerialHydration(ref: object): boolean { + if (isLargeValueRef(ref)) { + return ref.size > MAX_INLINE_MATERIALIZATION_BYTES + } + if (isLargeArrayManifest(ref)) { + return ref.chunks.some((chunk) => chunk.byteSize > MAX_INLINE_MATERIALIZATION_BYTES) + } + return false +} + /** * Dedupe the collected refs by identity, then replace each in parallel (bounded by - * {@link REF_CONCURRENCY}). `Map.set` is synchronous, so concurrent workers writing - * the shared map do not race. + * {@link REF_CONCURRENCY}); refs needing an oversized hydration run serially after + * the pool drains (see {@link requiresSerialHydration}). `Map.set` is synchronous, + * so concurrent workers writing the shared map do not race. * * `mapWithConcurrency`'s `fn` must not reject (a rejection fails the pool * non-deterministically), so the mapper is total: it catches per-ref errors and @@ -124,15 +188,28 @@ async function resolveReplacements( options: RedactLargeValueRefsOptions ): Promise> { const unique = [...new Set(refs)] + const pooled = unique.filter((ref) => !requiresSerialHydration(ref)) + const oversized = unique.filter(requiresSerialHydration) const replacements = new Map() let firstError: unknown - await mapWithConcurrency(unique, REF_CONCURRENCY, async (ref) => { + const resolveOne = async (ref: object, opts: RedactLargeValueRefsOptions): Promise => { try { - replacements.set(ref, await replaceRef(ref, options)) + replacements.set(ref, await replaceRef(ref, opts)) } catch (error) { if (firstError === undefined) firstError = error } - }) + } + await mapWithConcurrency(pooled, REF_CONCURRENCY, (ref) => resolveOne(ref, options)) + const gate = options.oversizedGate + for (const ref of oversized) { + if (!gate || options.withinOversizedGate) { + // Already serialized by a gated ancestor (or no gate: direct internal + // call) — re-acquiring would deadlock on the ancestor's own hold. + await resolveOne(ref, options) + } else { + await runSerially(gate, () => resolveOne(ref, { ...options, withinOversizedGate: true })) + } + } if (firstError !== undefined) throw firstError return replacements } @@ -176,21 +253,32 @@ async function replaceRef(ref: object, options: RedactLargeValueRefsOptions): Pr } /** - * Mask a materialized large value and re-offload it: handle any nested refs - * first, then mask inline strings, then re-store. `redactObjectStrings` skips - * refs, so the nested re-stored refs are left intact while their siblings mask. + * Mask a materialized value: mask inline strings first (`redactObjectStrings` + * treats refs as opaque), then replace any nested refs with their own masked + * results. Strings-first matters — a nested ref whose masked value shrinks + * below the offload threshold comes back inline, and running the string pass + * after substitution would mask that already-masked content a second time + * (non-idempotent for custom patterns). */ -async function maskAndReStore( +async function maskMaterializedValue( value: unknown, options: RedactLargeValueRefsOptions ): Promise { - const nested = await redactValueRefs(value, options) - const masked = await redactObjectStrings(nested, { + const masked = await redactObjectStrings(value, { entityTypes: options.entityTypes, language: options.language, customPatterns: options.customPatterns, onFailure: options.onFailure ?? 'scrub', }) + return redactValueRefs(masked, options) +} + +/** Mask a materialized large value and re-offload it as a fresh durable ref. */ +async function maskAndReStore( + value: unknown, + options: RedactLargeValueRefsOptions +): Promise { + const masked = await maskMaterializedValue(value, options) return compactExecutionPayload(masked, { ...options.store, requireDurable: true }) } @@ -217,6 +305,11 @@ async function redactRef( const materialized = await materializeLargeValueRef(ref, { ...options.store, trackReference: false, + // Redaction must hydrate refs past the default inline ceiling to mask + // them; refs above that ceiling are scheduled serially (see + // isOversizedSingleRef) so the raised budget never multiplies across the + // concurrency pool. + maxBytes: MAX_DURABLE_LARGE_VALUE_BYTES, }) if (materialized === undefined) { return onRefFailure( @@ -231,13 +324,47 @@ async function redactRef( } } +/** + * Mask a large-array manifest chunk-by-chunk: page one stored chunk at a time, + * mask nested refs + strings in that slice, and append the masked items to a + * fresh manifest. Peak memory stays ~one chunk regardless of the manifest's + * total `byteSize`, so a large offloaded array never trips the inline + * materialization ceiling. The rebuilt manifest re-chunks by byte target, + * recomputes `count`/`byteSize` bookkeeping from the masked items, and derives + * `preview` from masked content — the source manifest's preview holds raw items + * and must never be carried forward. + */ async function redactManifest( manifest: LargeArrayManifest, options: RedactLargeValueRefsOptions ): Promise { try { - const materialized = await materializeLargeArrayManifest(manifest, { ...options.store }) - return await maskAndReStore(materialized, options) + const readContext = { + ...options.store, + trackReference: false, + // Chunks target ~8MB, but a single item larger than the target still + // occupies its own chunk — allow those up to the durable cap since they + // hydrate one at a time. + maxBytes: MAX_DURABLE_LARGE_VALUE_BYTES, + } + let masked: LargeArrayManifest | undefined + let cursor = 0 + for (const chunk of manifest.chunks) { + const slice = await readLargeArrayManifestSlice(manifest, cursor, chunk.count, readContext) + cursor += chunk.count + const items = (await maskMaterializedValue(slice, options)) as unknown[] + masked = + masked === undefined + ? await createLargeArrayManifest(items, { ...options.store }) + : await appendLargeArrayManifest(masked, items, { ...options.store }) + } + if (masked === undefined) { + return createLargeArrayManifest([], { ...options.store }) + } + if (masked.totalCount !== manifest.totalCount) { + throw new Error('Masked manifest item count does not match the source manifest.') + } + return masked } catch (error) { return onRefFailure(error, options, 'Failed to redact large array manifest') } diff --git a/apps/sim/lib/logs/execution/pii-redaction.ts b/apps/sim/lib/logs/execution/pii-redaction.ts index d29c4cb864b..cfca0c8ea0e 100644 --- a/apps/sim/lib/logs/execution/pii-redaction.ts +++ b/apps/sim/lib/logs/execution/pii-redaction.ts @@ -145,7 +145,9 @@ function transformUnit( * * There is no total-size ceiling — the batching layer chunks the request and * fans out with bounded concurrency, so payloads of any size are masked properly - * rather than scrubbed. The per-chunk request timeout is the real backstop. + * rather than scrubbed. Transient chunk failures retry with backoff inside the + * mask client; a failure surfacing here means the retry budget is exhausted or + * the failure is deterministic. */ async function maskCollected( collected: string[],