Worker: ensure that final_state gets redacted

packages/engine-multi/CHANGELOG.md

@@ -1,5 +1,11 @@
 # engine-multi
 
+## 1.13.2
+
+### Patch Changes
+
+- 5999cf9: Fix a bug where an oversized final run state could reach Lightning without being redacted
+
 ## 1.13.1
 
 ### Patch Changes

packages/engine-multi/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openfn/engine-multi",
-  "version": "1.13.1",
+  "version": "1.13.2",
   "description": "Multi-process runtime engine",
   "main": "dist/index.js",
   "type": "module",

packages/engine-multi/src/api/lifecycle.ts

@@ -48,7 +48,7 @@ export const workflowComplete = (
   event: internalEvents.WorkflowCompleteEvent
 ) => {
   const { logger, state } = context;
-  const { workflowId, state: result, threadId } = event;
+  const { workflowId, state: result, threadId, redacted } = event;
 
   logger.success('complete workflow ', workflowId);
   state.status = 'done';
@@ -59,6 +59,7 @@ export const workflowComplete = (
     threadId,
     duration: state.duration,
     state: result,
+    redacted,
     time: timestamp(),
   });
 };

packages/engine-multi/src/events.ts

@@ -68,6 +68,7 @@ export interface WorkflowCompletePayload extends ExternalEvent {
   state: any;
   duration: number;
   time: bigint;
+  redacted?: boolean;
 }
 
 export interface WorkflowErrorPayload extends ExternalEvent {

packages/engine-multi/src/util/ensure-payload-size.ts

@@ -86,7 +86,14 @@ export default async (
         newPayload.payloadSize_b = sizeBytes;
       }
     } catch (e: any) {
-      Object.assign(newPayload[key], replacements[key] ?? replacements.default);
+      const replacement = replacements[key];
+      if (replacement) {
+        // A key-specific replacement (eg 'log') has a known, fixed shape -
+        // merge so other fields on it (time, level, ...) survive
+        Object.assign(newPayload[key], replacement);
+      } else {
+        newPayload[key] = replacements.default;
+      }
       newPayload.redacted = true;
       if (key === 'state') {
         newPayload.payloadSize_b = e.sizeBytes;

packages/engine-multi/src/worker/events.ts

@@ -49,6 +49,7 @@ export interface WorkflowStartEvent extends InternalEvent {}
 
 export interface WorkflowCompleteEvent extends InternalEvent {
   state: any;
+  redacted?: boolean;
 }
 
 export interface JobStartEvent extends InternalEvent {

packages/engine-multi/src/worker/thread/runtime.ts

@@ -62,8 +62,8 @@ export const publish = async (
     type,
     threadId,
     processId,
-    ...safePayload,
     ...payload,
+    ...safePayload,
   });
 };
 

packages/engine-multi/test/api/lifecycle.test.ts

@@ -124,6 +124,37 @@ test('workflowComplete: updates state', (t) => {
   t.assert(state.duration! > 0);
 });
 
+test('workflowComplete: forwards redacted', (t) => {
+  return new Promise((done) => {
+    const workflowId = 'a';
+
+    const state = {
+      id: workflowId,
+      startTime: Date.now() - 1000,
+    } as WorkflowState;
+    const context = createContext(workflowId, state);
+
+    const event: w.WorkflowCompleteEvent = {
+      type: w.WORKFLOW_COMPLETE,
+      workflowId,
+      state: { data: '[REDACTED]' },
+      threadId: '1',
+      redacted: true,
+    };
+
+    // Without this, ws-worker's run-complete handler has no way to know the
+    // final state was too big and got redacted - it just silently ships
+    // '[REDACTED]' with no explanation, unlike step-complete's handling of
+    // an oversized dataclip
+    context.on(e.WORKFLOW_COMPLETE, (evt) => {
+      t.true(evt.redacted);
+      done();
+    });
+
+    workflowComplete(context, event);
+  });
+});
+
 test(`job-start: emits ${e.JOB_START} with key fields`, (t) => {
   return new Promise((done) => {
     const workflowId = 'a';

packages/engine-multi/test/integration.test.ts

@@ -519,6 +519,7 @@ test.serial('redact final state if it exceeds the payload limit', (t) => {
     const expression = `
 export default [(state) => {
   state.data = new Array(1024 * 512).fill('a').join('')
+  state.otherStuff = 1001;
   return state;
 }]`;
 
@@ -535,7 +536,7 @@ export default [(state) => {
       .execute(plan, emptyState, options)
       .on('workflow-complete', ({ state }) => {
         t.log(state);
-        t.is(state.data, '[REDACTED]');
+        t.deepEqual(state, { data: '[REDACTED]' });
         done();
       });
   });
@@ -698,3 +699,43 @@ export default [(state) => {
     });
   }
 );
+
+test.serial(
+  'redact multi-leaf final state when only the combined size exceeds the limit',
+  (t) => {
+    return new Promise(async (done) => {
+      api = await createAPI({
+        logger,
+      });
+
+      // Each leaf on its own is well under the 0.3mb limit, so no per-job
+      // redaction fires - only the aggregated multi-leaf dict is oversized
+      const leafExpression = (n: number) => `${withFn}fn((state) => {
+  state.data = new Array(1024 * 150).fill('${n}').join('');
+  return state;
+})`;
+
+      const jobs = [
+        {
+          id: 'a',
+          next: { b: true, c: true, d: true },
+        },
+        { id: 'b', expression: leafExpression(1) },
+        { id: 'c', expression: leafExpression(2) },
+        { id: 'd', expression: leafExpression(3) },
+      ];
+
+      const plan = createPlan(jobs);
+      const options = {
+        payloadLimitMb: 0.3,
+      };
+
+      api
+        .execute(plan, emptyState, options)
+        .on('workflow-complete', ({ state }) => {
+          t.deepEqual(state, { data: '[REDACTED]' });
+          done();
+        });
+    });
+  }
+);

packages/lightning-mock/CHANGELOG.md

@@ -1,5 +1,12 @@
 # @openfn/lightning-mock
 
+## 2.4.29
+
+### Patch Changes
+
+- Updated dependencies [5999cf9]
+  - @openfn/engine-multi@1.13.2
+
 ## 2.4.28
 
 ### Patch Changes

packages/lightning-mock/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openfn/lightning-mock",
-  "version": "2.4.28",
+  "version": "2.4.29",
   "private": true,
   "description": "A mock Lightning server",
   "main": "dist/index.js",

packages/ws-worker/CHANGELOG.md

@@ -1,5 +1,13 @@
 # ws-worker
 
+## 1.29.3
+
+### Patch Changes
+
+- 5999cf9: Warn when a run's final state has been redacted for exceeding the payload size limit
+- Updated dependencies [5999cf9]
+  - @openfn/engine-multi@1.13.2
+
 ## 1.29.2
 
 ### Patch Changes

packages/ws-worker/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openfn/ws-worker",
-  "version": "1.29.2",
+  "version": "1.29.3",
   "description": "A Websocket Worker to connect Lightning to a Runtime Engine",
   "main": "dist/index.js",
   "type": "module",

packages/ws-worker/src/events/run-complete.ts

@@ -1,12 +1,14 @@
 import type { WorkflowCompletePayload } from '@openfn/engine-multi';
 import type { RunCompletePayload } from '@openfn/lexicon/lightning';
+import { timestamp } from '@openfn/logger';
 
 import { RUN_COMPLETE } from '../events';
 import { calculateRunExitReason } from '../api/reasons';
 import { Context } from '../api/execute';
 import logFinalReason from '../util/log-final-reason';
 import { timeInMicroseconds } from '../util';
 import { sendEvent } from '../util/send-event';
+import handleJobLog from './run-log';
 
 const isEmptyState = (obj: any) => {
   if (
@@ -71,6 +73,20 @@ export default async function onWorkflowComplete(
     ...reason,
   };
 
+  if (event.redacted) {
+    const time = (timestamp() - BigInt(10e6)).toString();
+    await handleJobLog(context, [
+      {
+        time,
+        message: [
+          'WARNING: Final state exceeds dataclip size limit. The dataclip has been redacted. If this is a cron workflow, the next run will be passed an invalid state object',
+        ],
+        level: 'info',
+        name: 'R/T',
+      },
+    ]);
+  }
+
   if (isSingleLeaf) {
     payload.final_dataclip_id = state.leafDataclipIds[0];
   }

packages/ws-worker/src/util/cli.ts

@@ -232,7 +232,7 @@ export default function parseArgs(argv: string[]): Args {
     })
     .option('payload-memory', {
       description:
-        'Maximum memory allocated to a single run, in mb. Env: WORKER_MAX_PAYLOAD_MB',
+        'Maximum serialized size of any payload, in mb. Env: WORKER_MAX_PAYLOAD_MB',
       type: 'number',
     })
     .option('stringify-state', {