perf(pii): mask offloaded refs + parallelize + raise redaction concurrency

apps/pii/server.py

@@ -181,6 +181,24 @@ class AnonymizeBatchRequest(BaseModel):
     operators: dict[str, dict[str, Any]] | None = None
 
 
+class RedactRequest(BaseModel):
+    text: str
+    language: str = "en"
+    entities: list[str] | None = None
+    score_threshold: float | None = None
+    anonymizers: dict[str, dict[str, Any]] | None = None
+    operators: dict[str, dict[str, Any]] | None = None
+
+
+class RedactBatchRequest(BaseModel):
+    texts: list[str]
+    language: str = "en"
+    entities: list[str] | None = None
+    score_threshold: float | None = None
+    anonymizers: dict[str, dict[str, Any]] | None = None
+    operators: dict[str, dict[str, Any]] | None = None
+
+
 def build_operators(
     raw_operators: dict[str, dict[str, Any]] | None,
 ) -> dict[str, OperatorConfig] | None:
@@ -296,3 +314,74 @@ def anonymize_batch(req: AnonymizeBatchRequest) -> dict[str, list[str]]:
             for item in req.items
         ]
     }
+
+
+@app.post("/redact")
+def redact(req: RedactRequest) -> dict[str, str]:
+    """Analyze + anonymize one text in a single round-trip (the combined
+    counterpart to /analyze followed by /anonymize). Returns masked text; a text
+    with no detected PII passes through unchanged. The analyzer results feed the
+    anonymizer directly (no dict round-trip)."""
+    started = time.perf_counter()
+    operators = build_operators(req.anonymizers or req.operators)
+    results = analyzer.analyze(
+        text=req.text,
+        language=req.language,
+        entities=req.entities or None,
+        score_threshold=req.score_threshold,
+    )
+    text = (
+        req.text
+        if not results
+        else anonymizer.anonymize(
+            text=req.text, analyzer_results=results, operators=operators
+        ).text
+    )
+    logger.info(
+        "redact lang=%s chars=%d spans=%d duration_ms=%.1f",
+        req.language,
+        len(req.text),
+        len(results),
+        (time.perf_counter() - started) * 1000,
+    )
+    return {"text": text}
+
+
+@app.post("/redact_batch")
+def redact_batch(req: RedactBatchRequest) -> dict[str, list[str]]:
+    """Analyze + anonymize many texts in a single round-trip (the combined
+    counterpart to /analyze_batch followed by /anonymize_batch). Returns masked
+    text per input in request order; texts with no detected PII pass through
+    unchanged. Analysis batches through spaCy nlp.pipe; the analyzer results feed
+    the anonymizer directly (no dict round-trip), and anonymization runs only on
+    texts that actually matched."""
+    started = time.perf_counter()
+    operators = build_operators(req.anonymizers or req.operators)
+    analyzed = list(
+        batch_analyzer.analyze_iterator(
+            texts=req.texts,
+            language=req.language,
+            entities=req.entities or None,
+            score_threshold=req.score_threshold,
+        )
+    )
+    masked: list[str] = []
+    total_spans = 0
+    for text, per_text in zip(req.texts, analyzed):
+        if not per_text:
+            masked.append(text)
+            continue
+        total_spans += len(per_text)
+        masked.append(
+            anonymizer.anonymize(
+                text=text, analyzer_results=per_text, operators=operators
+            ).text
+        )
+    logger.info(
+        "redact_batch lang=%s texts=%d spans=%d duration_ms=%.1f",
+        req.language,
+        len(req.texts),
+        total_spans,
+        (time.perf_counter() - started) * 1000,
+    )
+    return {"texts": masked}

apps/sim/executor/execution/block-executor.ts

@@ -4,6 +4,7 @@ import { redactApiKeys } from '@/lib/core/security/redaction'
 import { normalizeStringArray } from '@/lib/core/utils/arrays'
 import { getBaseUrl } from '@/lib/core/utils/urls'
 import { compactExecutionPayload } from '@/lib/execution/payloads/serializer'
+import { redactLargeValueRefsInValue } from '@/lib/logs/execution/pii-large-values'
 import { redactObjectStrings } from '@/lib/logs/execution/pii-redaction'
 import {
   containsUserFileWithMetadata,
@@ -228,15 +229,29 @@ export class BlockExecutor {
       }
 
       if (ctx.piiBlockOutputRedaction?.enabled) {
-        // In-flight redaction: mask before compaction (so offloaded large values
-        // are seen) and before the log/state split below, so both the downstream
-        // state copy and the persisted log copy are masked. `onFailure: 'throw'`
-        // aborts the run rather than feeding corrupted/leaked data downstream.
-        normalizedOutput = await redactObjectStrings(normalizedOutput, {
+        // In-flight redaction before the log/state split below, so both the
+        // downstream state copy and the persisted log copy are masked.
+        // `onFailure: 'throw'` aborts the run rather than feeding corrupted/leaked
+        // data downstream.
+        const redactionOptions = {
           entityTypes: ctx.piiBlockOutputRedaction.entityTypes,
           language: ctx.piiBlockOutputRedaction.language,
-          onFailure: 'throw',
+          onFailure: 'throw' as const,
+        }
+        // Tools like the function executor offload large outputs to large-value
+        // refs BEFORE they reach here, and the string walk treats a ref as opaque.
+        // So hydrate → mask → re-store any refs first, then mask inline strings —
+        // otherwise PII inside an offloaded output is never redacted.
+        normalizedOutput = await redactLargeValueRefsInValue(normalizedOutput, {
+          ...redactionOptions,
+          store: {
+            workspaceId: ctx.workspaceId,
+            workflowId: ctx.workflowId,
+            executionId: ctx.executionId,
+            userId: ctx.userId,
+          },
         })
+        normalizedOutput = await redactObjectStrings(normalizedOutput, redactionOptions)
       }
 
       normalizedOutput = (await compactExecutionPayload(normalizedOutput, {

apps/sim/lib/core/config/env.ts

@@ -328,7 +328,9 @@ export const env = createEnv({
     INTERNAL_API_BASE_URL:                 z.string().optional(),                  // Optional internal base URL for server-side self-calls; must include protocol if set (e.g., http://sim-app.namespace.svc.cluster.local:3000)
     ALLOWED_ORIGINS:                       z.string().optional(),                  // CORS allowed origins
     PII_URL:                               z.string().optional(),                  // Presidio PII service base URL serving /analyze + /anonymize (standalone ECS service; default http://localhost:5001 for local dev)
-    PII_MASK_CHUNK_CONCURRENCY:            z.coerce.number().int().positive().optional(), // Max in-flight mask-batch requests per redaction (default 4); raise for a scaled-out Presidio service, lower to 1 for a single instance
+    PII_MASK_CHUNK_CONCURRENCY:            z.coerce.number().int().positive().optional(), // Max in-flight mask-batch requests per redaction (default 64); tune to the Presidio fleet size behind the internal ALB, lower to 1 for a single instance
+    PII_REF_CONCURRENCY:                   z.coerce.number().int().positive().optional(), // Max large-value refs hydrated+masked+re-stored in parallel per payload (default 4); multiplies with PII_MASK_CHUNK_CONCURRENCY for total in-flight Presidio load
+    PII_SERVICE_CHUNK_CONCURRENCY:         z.coerce.number().int().positive().optional(), // Max Presidio requests in flight from a single mask-batch call (route -> Presidio fan-out, default 4); inner to PII_MASK_CHUNK_CONCURRENCY
 
     // OAuth Integration Credentials - All optional, enables third-party integrations
     GOOGLE_CLIENT_ID:                      z.string().optional(),                  // Google OAuth client ID for Google services

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

@@ -7,13 +7,15 @@ import { chunkIndicesByBudget } from '@/lib/guardrails/pii-batching'
 
 /**
  * Max in-flight mask-batch requests per call. Each request is a CPU-heavy NER
- * batch, so a single Presidio instance is easily saturated — default 4, raise it
- * via `PII_MASK_CHUNK_CONCURRENCY` for a scaled-out/load-balanced service, or set
- * 1 for a single instance. No request timeout: masking a large batch is slow and
- * the (scaled) Presidio service is expected to eventually respond; an unreachable
+ * batch — default 64, sized to saturate the load-balanced Presidio fleet behind
+ * the internal ALB (which spreads each request across tasks). Effective throughput
+ * is capped by fleet worker capacity, so past that this just queues; tune via
+ * `PII_MASK_CHUNK_CONCURRENCY` to the fleet size (and lower to 1 for a single
+ * self-hosted instance). No request timeout: masking a large batch is slow and the
+ * (scaled) Presidio service is expected to eventually respond; an unreachable
  * service still rejects fast (connection refused) so the caller scrubs.
  */
-const CHUNK_CONCURRENCY = env.PII_MASK_CHUNK_CONCURRENCY ?? 4
+const CHUNK_CONCURRENCY = env.PII_MASK_CHUNK_CONCURRENCY ?? 64
 
 /**
  * Mask PII across many strings via the internal app-container endpoint.

apps/sim/lib/guardrails/validate_pii.test.ts

@@ -35,6 +35,15 @@ describe('validate_pii (Presidio service)', () => {
     analyzeBodies = []
     fetchMock = vi.fn(async (url: string, init: { body: string }) => {
       const body = JSON.parse(init.body)
+      if (url.includes('/redact_batch')) {
+        for (const text of body.texts as string[]) {
+          analyzeBodies.push({ text, language: body.language, entities: body.entities })
+        }
+        const texts = (body.texts as string[]).map((t) =>
+          applyReplace(t, emailSpans(t, body.entities))
+        )
+        return new Response(JSON.stringify({ texts }), { status: 200 })
+      }
       if (url.includes('/analyze_batch')) {
         for (const text of body.texts as string[]) {
           analyzeBodies.push({ text, language: body.language, entities: body.entities })
@@ -89,7 +98,29 @@ describe('validate_pii (Presidio service)', () => {
 
     it('throws on a service failure so the caller can scrub', async () => {
       fetchMock.mockResolvedValueOnce(new Response('boom', { status: 500 }))
-      await expect(maskPIIBatch(['email a@b.com'], [])).rejects.toThrow(/Presidio analyze failed/)
+      await expect(maskPIIBatch(['email a@b.com'], [])).rejects.toThrow(
+        /Presidio redact_batch failed/
+      )
+    })
+
+    // Runs last: a 404 permanently flips the module's combined-endpoint flag off.
+    it('falls back to legacy analyze+anonymize when /redact_batch is absent (404)', async () => {
+      fetchMock.mockImplementation(async (url: string, init: { body: string }) => {
+        const body = JSON.parse(init.body)
+        if (url.includes('/redact_batch')) return new Response('Not Found', { status: 404 })
+        if (url.includes('/analyze_batch')) {
+          const spans = (body.texts as string[]).map((t) => emailSpans(t, body.entities))
+          return new Response(JSON.stringify(spans), { status: 200 })
+        }
+        // /anonymize_batch
+        const texts = (body.items as Array<{ text: string; analyzer_results: Span[] }>).map((i) =>
+          applyReplace(i.text, i.analyzer_results)
+        )
+        return new Response(JSON.stringify({ texts }), { status: 200 })
+      })
+
+      const out = await maskPIIBatch(['email a@b.com', 'clean'], [])
+      expect(out).toEqual(['email <EMAIL_ADDRESS>', 'clean'])
     })
   })
 

apps/sim/lib/guardrails/validate_pii.ts

@@ -7,14 +7,14 @@ import { chunkIndicesByBudget } from '@/lib/guardrails/pii-batching'
 const logger = createLogger('PIIValidator')
 
 /**
- * Concurrent chunk requests in flight. Each chunk is itself a batched service call
- * (spaCy `nlp.pipe` over many strings), so a small concurrency keeps a single-model
- * Presidio instance from holding too many parallel docs in memory while still
- * overlapping HTTP/JSON with the next chunk's NER.
+ * Concurrent chunk requests in flight from a single mask-batch call. Each chunk is
+ * itself a batched service call (spaCy `nlp.pipe` over many strings). Default 4;
+ * raise via `PII_SERVICE_CHUNK_CONCURRENCY` for a scaled Presidio fleet (this is
+ * the route → Presidio fan-out, inner to the app → route `PII_MASK_CHUNK_CONCURRENCY`).
  */
-const CHUNK_CONCURRENCY = 4
+const CHUNK_CONCURRENCY = env.PII_SERVICE_CHUNK_CONCURRENCY ?? 4
 
-/** Presidio service serving both /analyze and /anonymize (VIN is native there). */
+/** Presidio service serving /analyze, /anonymize, and combined /redact (VIN is native there). */
 const PII_URL = env.PII_URL || 'http://localhost:5001'
 
 export interface PIIValidationInput {
@@ -124,6 +124,50 @@ async function anonymizeBatch(items: AnonymizeBatchItem[]): Promise<string[]> {
   return data.texts
 }
 
+/**
+ * Flips to `false` the first time `/redact_batch` returns 404 — an older Presidio
+ * image without the combined endpoint — so subsequent chunks skip straight to the
+ * legacy analyze+anonymize path instead of re-probing. Reset on process restart
+ * (a deploy), so a newly-rolled Presidio is picked up.
+ */
+let combinedRedactAvailable = true
+
+/**
+ * Analyze + anonymize a batch in ONE round-trip via `/redact_batch`, returning
+ * masked texts in request order. Halves the app↔service round-trips (and avoids
+ * shipping the text back up for anonymize) vs {@link analyzeBatch} + {@link anonymizeBatch}.
+ *
+ * Returns `null` when the endpoint is absent (404 — an older Presidio image), so
+ * the caller falls back to the legacy two-call path. Any other non-2xx or a
+ * length mismatch throws so the caller applies its fail-safe (never leaks).
+ */
+async function redactBatch(
+  texts: string[],
+  entityTypes: string[],
+  language: string
+): Promise<string[] | null> {
+  const entities = entityTypes.length > 0 ? entityTypes : undefined
+
+  // boundary-raw-fetch: internal call to the Presidio combined redact service via PII_URL
+  const response = await fetch(`${PII_URL}/redact_batch`, {
+    method: 'POST',
+    headers: { 'content-type': 'application/json' },
+    body: JSON.stringify({ texts, language, ...(entities ? { entities } : {}) }),
+  })
+  if (response.status === 404) return null
+  if (!response.ok) {
+    const detail = await response.text().catch(() => '')
+    throw new Error(`Presidio redact_batch failed (${response.status}): ${detail.slice(0, 200)}`)
+  }
+  const data = (await response.json()) as { texts: string[] }
+  if (data.texts.length !== texts.length) {
+    throw new Error(
+      `Presidio redact_batch returned ${data.texts.length} result(s) for ${texts.length} input(s)`
+    )
+  }
+  return data.texts
+}
+
 /**
  * Mask spans via the Presidio anonymizer service. Omitting `anonymizers` uses the
  * default `replace` operator, which yields `<ENTITY_TYPE>`. Throws on failure.
@@ -212,12 +256,12 @@ export async function validatePII(input: PIIValidationInput): Promise<PIIValidat
  * Mask PII across many strings via the Presidio service, preserving input order.
  *
  * Strings are grouped into byte/count-budgeted chunks (see {@link chunkIndicesByBudget}),
- * and each chunk runs one batched `analyze` pass followed by one batched `anonymize`
- * pass over only the strings that actually matched — so the service round-trip count
- * scales with payload size, not leaf count, and spaCy batches NER via `nlp.pipe`.
- * Chunks run with bounded concurrency. Strings with no detected PII pass through
- * unchanged. Rejects on any service failure (which fails the whole batch) so callers
- * can apply their own fail-safe (scrub).
+ * and each chunk is masked via the combined `/redact_batch` endpoint (one analyze +
+ * anonymize round-trip). Against an older Presidio without that endpoint, it falls
+ * back to the legacy `analyze_batch` + `anonymize_batch` pair — so the app is safe
+ * to deploy before or after the service. Chunks run with bounded concurrency.
+ * Strings with no detected PII pass through unchanged. Rejects on any service
+ * failure (which fails the whole batch) so callers can apply their own fail-safe (scrub).
  */
 export async function maskPIIBatch(
   texts: string[],
@@ -230,6 +274,19 @@ export async function maskPIIBatch(
 
   await mapWithConcurrency(chunkIndicesByBudget(texts), CHUNK_CONCURRENCY, async (indices) => {
     const chunkTexts = indices.map((i) => texts[i])
+
+    if (combinedRedactAvailable) {
+      const masked = await redactBatch(chunkTexts, entityTypes, language)
+      if (masked) {
+        indices.forEach((originalIndex, pos) => {
+          result[originalIndex] = masked[pos]
+        })
+        return
+      }
+      // 404: older Presidio image; stop probing and use the legacy path below.
+      combinedRedactAvailable = false
+    }
+
     const spansPerText = await analyzeBatch(chunkTexts, entityTypes, language)
 
     // A short/misaligned batch response would silently leave the unmatched

apps/sim/lib/logs/execution/pii-large-values.test.ts

@@ -88,6 +88,40 @@ describe('redactLargeValueRefs', () => {
     expect(result.finalOutput).toBe('[REDACTION_FAILED]')
   })
 
+  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' }
+    )
+
+    const result = await redactLargeValueRefs(
+      { finalOutput: refA, traceSpans: [{ output: refB }] },
+      { entityTypes: [], language: 'en', store: STORE }
+    )
+
+    expect(result.finalOutput).toEqual({ note: 'MASKED(call bob)' })
+    expect((result.traceSpans as any[])[0].output).toEqual({ note: 'MASKED(email amy)' })
+    expect(mockMaterializeRef).toHaveBeenCalledTimes(2)
+    expect(mockCompact).toHaveBeenCalledTimes(2)
+  })
+
+  it('aborts (throws) when one of several refs fails in throw mode', async () => {
+    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
+    )
+
+    await expect(
+      redactLargeValueRefs(
+        { finalOutput: refA, traceSpans: [{ output: refB }] },
+        { entityTypes: [], language: 'en', store: STORE, onFailure: 'throw' }
+      )
+    ).rejects.toBeInstanceOf(PiiRedactionError)
+  })
+
   it('leaves payloads without refs untouched', async () => {
     const payload = { finalOutput: { answer: 'world', count: 5 } }
     const result = await redactLargeValueRefs(payload, {