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

* perf(pii): mask offloaded refs in block outputs/input + parallelize + raise concurrency - Block-output and input stages now hydrate → mask → re-store large-value refs (function/tool outputs are offloaded to refs before reaching the redactor, which treats refs as opaque) — fixes big block outputs coming back unredacted - Parallelize large-value ref hydrate/mask/re-store: collect refs across the whole payload and process them with bounded concurrency instead of one sequential pass per key (PII_REF_CONCURRENCY, default 4) - Raise mask-batch chunk concurrency default 4 -> 64 to saturate the load-balanced Presidio fleet behind the internal ALB (PII_MASK_CHUNK_CONCURRENCY, env-tunable) * fix(pii): keep resolveReplacements mapper total (respect mapWithConcurrency contract) - Catch per-ref errors inside the mapWithConcurrency mapper and rethrow the first after the pool drains, instead of letting a throwing mapper reject the pool (the util documents fn MUST NOT reject). Preserves fail-fast abort in throw mode. - Add multi-ref throw-mode test: one of several refs failing aborts the redaction * feat(pii): make route->Presidio chunk concurrency env-tunable (PII_SERVICE_CHUNK_CONCURRENCY) Was hardcoded to 4; now env-tunable (default 4) so the inner route->Presidio fan-out can scale with the fleet alongside the outer PII_MASK_CHUNK_CONCURRENCY. * feat(pii): combined /redact + /redact_batch endpoint (one round-trip) - Presidio: add /redact and /redact_batch (analyze + anonymize server-side, feeding analyzer results straight to the anonymizer, no dict round-trip). All existing endpoints kept, so old clients keep working during rollout. - App: maskPIIBatch now uses /redact_batch (halves app<->service round-trips, no shipping text back up for anonymize). Falls back to the legacy analyze_batch + anonymize_batch pair on a 404 (older Presidio image), so the app is safe to deploy before or after the service in either order. Same length-check fail-closed guarantee.
2026-09-01 14:59:19 +08:00 · 2026-07-06 13:33:44 -07:00
parent 9538beb1b8
commit 34b2abadcb
9 changed files with 323 additions and 37 deletions
@@ -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}
@@ -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, {
@@ -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
@@ -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.
@@ -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'])
    })
  })

@@ -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
@@ -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, {
@@ -1,5 +1,7 @@
 import { createLogger } from '@sim/logger'
 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 { isLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest-metadata'
@@ -45,17 +47,27 @@ export interface RedactLargeValueRefsOptions {
 *
 * Traversal is SYNCHRONOUS (collect refs, then substitute) so a large ref-free
 * payload costs only a cheap walk — never a promise-per-node. Only the handful of
- * actual refs incur async hydrate → mask → re-store work.
+ * actual refs incur async hydrate → mask → re-store work, and those run in
+ * parallel with bounded concurrency ({@link REF_CONCURRENCY}).
 */
 export async function redactLargeValueRefs(
  payload: RedactablePayload,
  options: RedactLargeValueRefsOptions
 ): Promise<RedactablePayload> {
+  // 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.
+  const refs: object[] = []
+  const seen = new WeakSet<object>()
+  for (const key of Object.keys(payload) as (keyof RedactablePayload)[]) {
+    if (payload[key] !== undefined) collectRefs(payload[key], refs, seen)
+  }
+  if (refs.length === 0) return payload
+
+  const replacements = await resolveReplacements(refs, options)
  const result: RedactablePayload = { ...payload }
  for (const key of Object.keys(payload) as (keyof RedactablePayload)[]) {
-    if (payload[key] !== undefined) {
-      result[key] = await redactValueRefs(payload[key], options)
-    }
+    if (payload[key] !== undefined) result[key] = substituteRefs(payload[key], replacements)
  }
  return result
 }
@@ -82,14 +94,46 @@ async function redactValueRefs(
  collectRefs(value, refs, new WeakSet())
  if (refs.length === 0) return value

-  const replacements = new Map<object, unknown>()
-  for (const ref of refs) {
-    if (replacements.has(ref)) continue
-    replacements.set(ref, await replaceRef(ref, options))
-  }
+  const replacements = await resolveReplacements(refs, options)
  return substituteRefs(value, replacements)
 }

+/**
+ * Max large-value refs hydrated → masked → re-stored in parallel per payload.
+ * Multiplies with the mask-batch chunk concurrency for total in-flight Presidio
+ * load, which the load-balanced fleet behind the internal ALB absorbs.
+ */
+const REF_CONCURRENCY = env.PII_REF_CONCURRENCY ?? 4
+
+/**
+ * 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.
+ *
+ * `mapWithConcurrency`'s `fn` must not reject (a rejection fails the pool
+ * non-deterministically), so the mapper is total: it catches per-ref errors and
+ * records the first one. In `onFailure: 'throw'` mode `replaceRef` throws, so after
+ * the pool drains we rethrow that first error — an unmaskable ref still aborts the
+ * run rather than passing through. In `'scrub'` mode `replaceRef` never throws.
+ */
+async function resolveReplacements(
+  refs: object[],
+  options: RedactLargeValueRefsOptions
+): Promise<Map<object, unknown>> {
+  const unique = [...new Set(refs)]
+  const replacements = new Map<object, unknown>()
+  let firstError: unknown
+  await mapWithConcurrency(unique, REF_CONCURRENCY, async (ref) => {
+    try {
+      replacements.set(ref, await replaceRef(ref, options))
+    } catch (error) {
+      if (firstError === undefined) firstError = error
+    }
+  })
+  if (firstError !== undefined) throw firstError
+  return replacements
+}
+
 /** Depth-first sync walk collecting ref/manifest nodes (not recursing into them). */
 function collectRefs(value: unknown, out: object[], seen: WeakSet<object>): void {
  if (isLargeValueRef(value) || isLargeArrayManifest(value)) {
@@ -680,12 +680,24 @@ async function executeWorkflowCoreImpl(
    if (piiRedaction.input.enabled) {
      // Redact the input before the workflow sees it. `onFailure: 'throw'` aborts
      // the run (handled by the surrounding catch) rather than feeding a scrub
-      // marker into execution or leaking unredacted input.
-      processedInput = await redactObjectStrings(processedInput, {
+      // marker into execution or leaking unredacted input. A large input may
+      // already be offloaded to a large-value ref (opaque to the string walk), so
+      // hydrate → mask → re-store refs first, then mask inline strings.
+      const inputOpts = {
        entityTypes: piiRedaction.input.entityTypes,
        language: piiRedaction.input.language,
-        onFailure: 'throw',
+        onFailure: 'throw' as const,
+      }
+      processedInput = await redactLargeValueRefsInValue(processedInput, {
+        ...inputOpts,
+        store: {
+          workspaceId: providedWorkspaceId,
+          workflowId,
+          executionId,
+          userId: userId ?? undefined,
+        },
      })
+      processedInput = await redactObjectStrings(processedInput, inputOpts)
    }

    if (piiRedaction.blockOutputs.enabled) {