fix(data-retention): re-mask offloaded large-value refs on resume + don't lock out granular saves

- Resume/run-from-block restore now hydrates → masks → re-stores large-value refs
  in restored blockStates (not just inline strings), so a value offloaded before the
  block-output stage was enabled can't warm raw PII into downstream blocks. Fails fast.
- pii-large-values: add onFailure mode (throw on the execution path, scrub for logs)
  and redactLargeValueRefsInValue for arbitrary (non-RedactablePayload) values
- Granular flag gate now rejects only NEW off→on granular enablement, so orgs that
  already configured granular stages can still save retention settings when the flag is off

Author: TheodoreSpeaks

apps/sim/app/api/organizations/[id]/data-retention/route.ts

@@ -65,6 +65,27 @@ function normalizeConfigured(
   }
 }
 
+/**
+ * Which granular stages (`input`/`blockOutputs`) are already enabled per rule
+ * target (`workspaceId ?? ''` = the org default). Used to gate the
+ * `pii-granular-redaction` flag on *new* enablement only: when the flag is off,
+ * an org that already configured granular stages must still be able to re-save
+ * unrelated settings (the UI re-sends the full PII snapshot every save), so we
+ * reject only a stage transitioning off→on, never a preserved one.
+ */
+function granularStageEnablement(
+  settings: OrganizationRetentionValues['piiRedaction']
+): Map<string, { input: boolean; blockOutputs: boolean }> {
+  const map = new Map<string, { input: boolean; blockOutputs: boolean }>()
+  for (const rule of settings?.rules ?? []) {
+    map.set(rule.workspaceId ?? '', {
+      input: rule.stages?.input?.enabled === true,
+      blockOutputs: rule.stages?.blockOutputs?.enabled === true,
+    })
+  }
+  return map
+}
+
 /**
  * GET /api/organizations/[id]/data-retention
  * Returns the organization's data retention settings.
@@ -210,18 +231,28 @@ export const PUT = withRouteHandler(
           { status: 403 }
         )
       }
-      const enablesGranularStage = (body.piiRedaction?.rules ?? []).some(
-        (rule) =>
-          rule.stages?.input?.enabled === true || rule.stages?.blockOutputs?.enabled === true
-      )
-      if (!piiGranularRedactionEnabled && enablesGranularStage) {
-        return NextResponse.json(
-          {
-            error:
-              'Granular PII redaction (workflow input and block outputs) is not enabled for this organization',
-          },
-          { status: 403 }
-        )
+      if (!piiGranularRedactionEnabled) {
+        // Reject only a granular stage transitioning off→on; a body that merely
+        // preserves already-enabled granular stages must still save (the UI
+        // re-sends the full snapshot on every save), so existing orgs aren't
+        // locked out of unrelated retention changes when the flag is off.
+        const currentGranular = granularStageEnablement(current.piiRedaction)
+        const newlyEnablesGranular = (body.piiRedaction?.rules ?? []).some((rule) => {
+          const cur = currentGranular.get(rule.workspaceId ?? '')
+          return (
+            (rule.stages?.input?.enabled === true && !cur?.input) ||
+            (rule.stages?.blockOutputs?.enabled === true && !cur?.blockOutputs)
+          )
+        })
+        if (newlyEnablesGranular) {
+          return NextResponse.json(
+            {
+              error:
+                'Granular PII redaction (workflow input and block outputs) is not enabled for this organization',
+            },
+            { status: 403 }
+          )
+        }
       }
       merged.piiRedaction = body.piiRedaction
     }

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

@@ -28,7 +28,11 @@ vi.mock('@/lib/guardrails/mask-client', () => ({
   maskPIIBatchViaHttp: mockMaskBatch,
 }))
 
-import { redactLargeValueRefs } from '@/lib/logs/execution/pii-large-values'
+import {
+  redactLargeValueRefs,
+  redactLargeValueRefsInValue,
+} from '@/lib/logs/execution/pii-large-values'
+import { PiiRedactionError } from '@/lib/logs/execution/pii-redaction'
 
 const REF = {
   __simLargeValueRef: true,
@@ -95,3 +99,48 @@ describe('redactLargeValueRefs', () => {
     expect(mockMaterializeRef).not.toHaveBeenCalled()
   })
 })
+
+describe('redactLargeValueRefsInValue (arbitrary blockStates)', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockMaskBatch.mockImplementation(async (texts: string[]) => texts.map((t) => `MASKED(${t})`))
+    mockCompact.mockImplementation(async (value: unknown) => value)
+  })
+
+  it('hydrates + re-stores a ref nested in a non-RedactablePayload shape', async () => {
+    mockMaterializeRef.mockResolvedValue({ note: 'contact bob' })
+    const blockStates = { 'block-1': { output: REF }, 'block-2': { output: { plain: 'hi' } } }
+
+    const result = await redactLargeValueRefsInValue(blockStates, {
+      entityTypes: ['PERSON'],
+      language: 'en',
+      store: STORE,
+    })
+
+    expect((result as any)['block-1'].output).toEqual({ note: 'MASKED(contact bob)' })
+    expect((result as any)['block-2'].output).toEqual({ plain: 'hi' })
+  })
+
+  it('throws PiiRedactionError on failure when onFailure is throw (aborts resume, no marker)', async () => {
+    mockMaterializeRef.mockResolvedValue(undefined)
+
+    await expect(
+      redactLargeValueRefsInValue(
+        { 'block-1': { output: REF } },
+        { entityTypes: [], language: 'en', store: STORE, onFailure: 'throw' }
+      )
+    ).rejects.toBeInstanceOf(PiiRedactionError)
+  })
+
+  it('rethrows a re-store failure as PiiRedactionError under throw mode', async () => {
+    mockMaterializeRef.mockResolvedValue({ note: 'secret@x.com' })
+    mockCompact.mockRejectedValueOnce(new Error('s3 down'))
+
+    await expect(
+      redactLargeValueRefsInValue(
+        { 'block-1': { output: REF } },
+        { entityTypes: [], language: 'en', store: STORE, onFailure: 'throw' }
+      )
+    ).rejects.toBeInstanceOf(PiiRedactionError)
+  })
+})

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

@@ -8,6 +8,8 @@ import { compactExecutionPayload } from '@/lib/execution/payloads/serializer'
 import type { LargeValueStoreContext } from '@/lib/execution/payloads/store'
 import { materializeLargeValueRef } from '@/lib/execution/payloads/store'
 import {
+  PiiRedactionError,
+  type PiiRedactionFailureMode,
   REDACTION_FAILED_MARKER,
   type RedactablePayload,
   redactObjectStrings,
@@ -21,6 +23,13 @@ export interface RedactLargeValueRefsOptions {
   language: string
   /** Storage scope for materializing and re-storing the masked values. */
   store: LargeValueStoreContext
+  /**
+   * How to handle a ref that can't be materialized/masked/re-stored. Defaults to
+   * `'scrub'` (marker) — safe for the logs path. The execution-altering restore
+   * path passes `'throw'` so an unmaskable restored value aborts the run rather
+   * than feeding a marker (or leaking raw bytes) downstream.
+   */
+  onFailure?: PiiRedactionFailureMode
 }
 
 /**
@@ -51,6 +60,19 @@ export async function redactLargeValueRefs(
   return result
 }
 
+/**
+ * Hydrate, mask, and re-store offloaded large values inside an arbitrary value
+ * (e.g. resumed snapshot `blockStates`) — the same walk as
+ * {@link redactLargeValueRefs}, but not bound to the {@link RedactablePayload}
+ * key set. Structure is preserved; only refs/manifests are replaced.
+ */
+export async function redactLargeValueRefsInValue<T>(
+  value: T,
+  options: RedactLargeValueRefsOptions
+): Promise<T> {
+  return (await redactValueRefs(value, options)) as T
+}
+
 /** Sync-collect every ref/manifest in `value`, then async-replace each, then sync-substitute. */
 async function redactValueRefs(
   value: unknown,
@@ -119,11 +141,26 @@ async function maskAndReStore(
   const masked = await redactObjectStrings(nested, {
     entityTypes: options.entityTypes,
     language: options.language,
-    onFailure: 'scrub',
+    onFailure: options.onFailure ?? 'scrub',
   })
   return compactExecutionPayload(masked, { ...options.store, requireDurable: true })
 }
 
+/** Rethrow (as {@link PiiRedactionError}) or scrub-to-marker, per `onFailure`. */
+function onRefFailure(
+  error: unknown,
+  options: RedactLargeValueRefsOptions,
+  context: string
+): never | string {
+  if ((options.onFailure ?? 'scrub') === 'throw') {
+    throw error instanceof PiiRedactionError
+      ? error
+      : new PiiRedactionError(`${context}: ${getErrorMessage(error)}`)
+  }
+  logger.error(`${context}; scrubbing`, { error: getErrorMessage(error) })
+  return REDACTION_FAILED_MARKER
+}
+
 async function redactRef(
   ref: LargeValueRef,
   options: RedactLargeValueRefsOptions
@@ -133,11 +170,16 @@ async function redactRef(
       ...options.store,
       trackReference: false,
     })
-    if (materialized === undefined) return REDACTION_FAILED_MARKER
+    if (materialized === undefined) {
+      return onRefFailure(
+        new Error('large value could not be materialized'),
+        options,
+        'Failed to redact large value ref'
+      )
+    }
     return await maskAndReStore(materialized, options)
   } catch (error) {
-    logger.error('Failed to redact large value ref; scrubbing', { error: getErrorMessage(error) })
-    return REDACTION_FAILED_MARKER
+    return onRefFailure(error, options, 'Failed to redact large value ref')
   }
 }
 
@@ -149,9 +191,6 @@ async function redactManifest(
     const materialized = await materializeLargeArrayManifest(manifest, { ...options.store })
     return await maskAndReStore(materialized, options)
   } catch (error) {
-    logger.error('Failed to redact large array manifest; scrubbing', {
-      error: getErrorMessage(error),
-    })
-    return REDACTION_FAILED_MARKER
+    return onRefFailure(error, options, 'Failed to redact large array manifest')
   }
 }

apps/sim/lib/workflows/executor/execution-core.ts

@@ -18,6 +18,7 @@ import { clearExecutionCancellation } from '@/lib/execution/cancellation'
 import { warmLargeValueRefs } from '@/lib/execution/payloads/hydration'
 import { parseLargeExecutionValue } from '@/lib/execution/payloads/large-execution-value'
 import type { LoggingSession } from '@/lib/logs/execution/logging-session'
+import { redactLargeValueRefsInValue } from '@/lib/logs/execution/pii-large-values'
 import { redactObjectStrings } from '@/lib/logs/execution/pii-redaction'
 import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
 import {
@@ -675,25 +676,38 @@ export async function executeWorkflowCore(
       // predate the blockOutputs stage being enabled, re-mask them so downstream
       // blocks can't read unredacted PII from restored snapshot state. Masking is
       // idempotent, so outputs already masked in the original run are unaffected.
-      // Limitation: this walks inline strings only — values offloaded to
-      // large-value storage are still refs here and are not re-masked. In the
-      // normal flow that is safe (a run with the stage on masks before offload);
-      // the gap is the narrow case of a run that offloaded a large value while
-      // the stage was OFF and is resumed after the stage is turned ON.
+      //
+      // Two disjoint passes cover the whole state: `redactLargeValueRefsInValue`
+      // hydrates → masks → re-stores any value offloaded to large-value storage
+      // (>8MB refs the string walk treats as opaque), then `redactObjectStrings`
+      // masks the remaining inline string leaves. Both fail-fast (`throw`), so an
+      // unmaskable restored value aborts the resume rather than warming raw PII
+      // into `blockStates` for downstream blocks.
       const blockOutputOpts = {
         entityTypes: piiRedaction.blockOutputs.entityTypes,
         language: piiRedaction.blockOutputs.language,
         onFailure: 'throw' as const,
       }
+      const largeRefOpts = {
+        ...blockOutputOpts,
+        store: {
+          workspaceId: providedWorkspaceId,
+          workflowId,
+          executionId,
+          userId: userId ?? undefined,
+        },
+      }
       if (snapshot.state?.blockStates) {
-        snapshot.state.blockStates = await redactObjectStrings(
-          snapshot.state.blockStates,
-          blockOutputOpts
-        )
+        const hydrated = await redactLargeValueRefsInValue(snapshot.state.blockStates, largeRefOpts)
+        snapshot.state.blockStates = await redactObjectStrings(hydrated, blockOutputOpts)
       }
       if (runFromBlock?.sourceSnapshot?.blockStates) {
-        runFromBlock.sourceSnapshot.blockStates = await redactObjectStrings(
+        const hydrated = await redactLargeValueRefsInValue(
           runFromBlock.sourceSnapshot.blockStates,
+          largeRefOpts
+        )
+        runFromBlock.sourceSnapshot.blockStates = await redactObjectStrings(
+          hydrated,
           blockOutputOpts
         )
       }