wip

Author: taylordowns2000

packages/lexicon/lightning.d.ts

@@ -169,7 +169,8 @@ export type RunStartReply = {}; // no payload
 
 export type RunCompletePayload = ExitReason & {
   timestamp: TimeInMicroSeconds;
-  final_state?: any; // The aggregated final state from the workflow (handles branching)
+  final_dataclip_id?: string; // Single-leaf: reuse the step's output dataclip
+  final_state?: any; // Multi-leaf: aggregated final state keyed by step id
 };
 export type RunCompleteReply = undefined;
 

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

@@ -14,19 +14,28 @@ export default async function onWorkflowComplete(
 ) {
   const { state, onFinish, logger } = context;
 
-  // Use the aggregated final state from the runtime
-  // This handles branching workflows correctly by returning all leaf states
   const result = event.state;
 
   const reason = calculateRunExitReason(state);
   await logFinalReason(context, reason);
 
+  const isSingleLeaf =
+    state.leafDataclipIds.length === 1 &&
+    !state.withheldDataclips[state.leafDataclipIds[0]];
+
+  const payload: RunCompletePayload = {
+    timestamp: timeInMicroseconds(event.time),
+    ...reason,
+  };
+
+  if (isSingleLeaf) {
+    payload.final_dataclip_id = state.leafDataclipIds[0];
+  } else {
+    payload.final_state = result;
+  }
+
   try {
-    await sendEvent<RunCompletePayload>(context, RUN_COMPLETE, {
-      final_state: result,
-      timestamp: timeInMicroseconds(event.time),
-      ...reason,
-    });
+    await sendEvent<RunCompletePayload>(context, RUN_COMPLETE, payload);
   } catch (e) {
     logger.error(
       `${state.plan.id} failed to send ${RUN_COMPLETE} event. This run will be lost!`

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

@@ -32,13 +32,10 @@ export default async function onStepComplete(
   delete state.activeStep;
   delete state.activeJob;
 
-  // TODO right now, the last job to run will be the result for the run
-  // this may not stand up in the future
-  // I'd feel happer if the runtime could judge what the final result is
-  // (taking into account branches and stuff)
-  // The problem is that the runtime will return the object, not an id,
-  // so we have a bit of a mapping problem
-  state.lastDataclipId = dataclipId;
+  // Track leaf dataclips (steps with no downstream jobs)
+  if (!event.next?.length) {
+    state.leafDataclipIds.push(dataclipId);
+  }
 
   // Set the input dataclip id for downstream jobs
   event.next?.forEach((nextJobId) => {

packages/ws-worker/src/types.d.ts

@@ -18,8 +18,8 @@ export type RunState = {
   withheldDataclips: Record<string, true>;
   reasons: Record<string, ExitReason>;
 
-  // final dataclip id
-  lastDataclipId?: string;
+  // dataclip ids for leaf nodes (steps with no downstream)
+  leafDataclipIds: string[];
 };
 
 export type CancelablePromise = Promise<void> & {

packages/ws-worker/src/util/create-run-state.ts

@@ -3,7 +3,7 @@ import type { RunState } from '../types';
 
 export default (plan: ExecutionPlan, input?: Lazy<State>): RunState => {
   const state = {
-    lastDataclipId: '',
+    leafDataclipIds: [],
     dataclips: {},
     inputDataclips: {},
     withheldDataclips: {},

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

@@ -8,17 +8,43 @@ import { RUN_COMPLETE, RUN_LOG } from '../../src/events';
 import { createRunState } from '../../src/util';
 import { createPlan } from '../util';
 
-test('should send a run:complete event', async (t) => {
+test('should send a run:complete event with final_dataclip_id for single leaf', async (t) => {
   const result = { answer: 42 };
   const plan = createPlan();
 
   const state = createRunState(plan);
+  state.leafDataclipIds = ['clip-1'];
+
+  const channel = mockChannel({
+    [RUN_LOG]: () => true,
+    [RUN_COMPLETE]: (evt) => {
+      t.is(evt.final_dataclip_id, 'clip-1');
+      t.falsy(evt.final_state);
+      t.falsy(evt.time);
+    },
+  });
+
+  const event: any = { state: result };
+
+  const context: any = { channel, state, onFinish: () => {} };
+  await handleRunComplete(context, event);
+});
+
+test('should send final_state when there are multiple leaves', async (t) => {
+  const result = {
+    'job-a': { data: { a: true } },
+    'job-b': { data: { b: true } },
+  };
+  const plan = createPlan();
+
+  const state = createRunState(plan);
+  state.leafDataclipIds = ['clip-1', 'clip-2'];
 
   const channel = mockChannel({
     [RUN_LOG]: () => true,
     [RUN_COMPLETE]: (evt) => {
       t.deepEqual(evt.final_state, result);
-      t.falsy(evt.time); // if no timestamp in the engine event, no timestamp in the worker one
+      t.falsy(evt.final_dataclip_id);
     },
   });
 
@@ -240,9 +266,11 @@ test('should call onFinish even if the lightning event timesout', async (t) => {
   await handleRunComplete(context, event);
 });
 
-test('should send final_state for a linear workflow', async (t) => {
+test('should send final_dataclip_id for a single-leaf workflow', async (t) => {
   const plan = createPlan();
   const state = createRunState(plan);
+  state.leafDataclipIds = ['abc-123'];
+
   const finalResult = { data: { count: 100 }, references: [] };
 
   let completeEvent: any;
@@ -264,15 +292,16 @@ test('should send final_state for a linear workflow', async (t) => {
 
   await handleRunComplete(context, event);
 
-  t.deepEqual(completeEvent.final_state, finalResult);
+  t.is(completeEvent.final_dataclip_id, 'abc-123');
+  t.falsy(completeEvent.final_state);
   t.is(completeEvent.reason, 'success');
 });
 
 test('should send final_state for a branching workflow with multiple leaf nodes', async (t) => {
   const plan = createPlan();
   const state = createRunState(plan);
+  state.leafDataclipIds = ['clip-1', 'clip-2', 'clip-3'];
 
-  // Simulate a branching workflow with multiple final states
   const branchedResult = {
     'job-1': { data: { path: 'A', value: 42 } },
     'job-2': { data: { path: 'B', value: 84 } },
@@ -292,7 +321,6 @@ test('should send final_state for a branching workflow with multiple leaf nodes'
     channel,
     state,
     onFinish: ({ state: finalState }: any) => {
-      // Verify that onFinish receives the branched result
       t.deepEqual(finalState, branchedResult);
     },
   };
@@ -301,33 +329,18 @@ test('should send final_state for a branching workflow with multiple leaf nodes'
 
   await handleRunComplete(context, event);
 
-  // Verify the event contains the full branched state structure
   t.deepEqual(completeEvent.final_state, branchedResult);
+  t.falsy(completeEvent.final_dataclip_id);
   t.is(completeEvent.reason, 'success');
-  t.truthy(completeEvent.final_state['job-1']);
-  t.truthy(completeEvent.final_state['job-2']);
-  t.truthy(completeEvent.final_state['job-3']);
 });
 
-test('should properly serialize final_state as JSON', async (t) => {
+test('should send final_state when single leaf dataclip was withheld', async (t) => {
   const plan = createPlan();
   const state = createRunState(plan);
+  state.leafDataclipIds = ['clip-1'];
+  state.withheldDataclips = { 'clip-1': true };
 
-  // Test with complex state including nested objects, arrays, and special values
-  const complexState = {
-    data: {
-      users: [
-        { id: 1, name: 'Alice' },
-        { id: 2, name: 'Bob' },
-      ],
-      metadata: {
-        timestamp: new Date('2024-01-01').toISOString(),
-        nested: { deeply: { value: 42 } },
-      },
-    },
-    configuration: { setting: true },
-    references: [],
-  };
+  const result = { data: { big: 'data' } };
 
   let completeEvent: any;
 
@@ -344,28 +357,23 @@ test('should properly serialize final_state as JSON', async (t) => {
     onFinish: () => {},
   };
 
-  const event: any = { state: complexState };
+  const event: any = { state: result };
 
   await handleRunComplete(context, event);
 
-  // Verify the state is properly preserved
-  t.deepEqual(completeEvent.final_state, complexState);
-  t.deepEqual(completeEvent.final_state.data.users[0], { id: 1, name: 'Alice' });
-  t.is(completeEvent.final_state.data.metadata.nested.deeply.value, 42);
-
-  // Verify it can be stringified (simulating what happens when sent over the wire)
-  const jsonString = JSON.stringify(completeEvent.final_state);
-  const parsed = JSON.parse(jsonString);
-  t.deepEqual(parsed, complexState);
+  t.deepEqual(completeEvent.final_state, result);
+  t.falsy(completeEvent.final_dataclip_id);
 });
 
-test('should handle Uint8Array in final_state', async (t) => {
+test('should send final_state when a single leaf node is reached by two paths', async (t) => {
   const plan = createPlan();
   const state = createRunState(plan);
+  // Same node executed twice via different paths produces two leaf dataclips
+  state.leafDataclipIds = ['clip-1', 'clip-2'];
 
-  // Test with Uint8Array which needs special handling
-  const stateWithBinary = {
-    data: { buffer: new Uint8Array([1, 2, 3, 4, 5]) },
+  const result = {
+    x: { data: { from: 'a' } },
+    'x-1': { data: { from: 'b' } },
   };
 
   let completeEvent: any;
@@ -383,10 +391,10 @@ test('should handle Uint8Array in final_state', async (t) => {
     onFinish: () => {},
   };
 
-  const event: any = { state: stateWithBinary };
+  const event: any = { state: result };
 
   await handleRunComplete(context, event);
 
-  // Verify the Uint8Array is preserved in the event
-  t.deepEqual(completeEvent.final_state.data.buffer, new Uint8Array([1, 2, 3, 4, 5]));
+  t.deepEqual(completeEvent.final_state, result);
+  t.falsy(completeEvent.final_dataclip_id);
 });

packages/ws-worker/test/events/run-error.test.ts

@@ -11,7 +11,7 @@ test('runError should trigger runComplete with a reason', async (t) => {
   const jobId = 'job-1';
 
   const state = createRunState(plan);
-  state.lastDataclipId = 'x';
+
   state.activeStep = 'b';
   state.activeJob = jobId;
 
@@ -40,7 +40,7 @@ test('workflow error should send reason to onFinish', async (t) => {
   const jobId = 'job-1';
 
   const state = createRunState(plan);
-  state.lastDataclipId = 'x';
+
   state.activeStep = 'b';
   state.activeJob = jobId;
 
@@ -72,7 +72,7 @@ test('workflow error should send reason to onFinish', async (t) => {
 
 test('runError should not call job complete if the job is not active', async (t) => {
   const state = createRunState(plan);
-  state.lastDataclipId = 'x';
+
 
   const channel = mockChannel({
     [RUN_LOG]: () => true,
@@ -113,7 +113,7 @@ test('runError should log the reason', async (t) => {
     },
     options: {},
   });
-  state.lastDataclipId = 'x';
+
   state.activeStep = 'b';
   state.activeJob = jobId;
 

packages/ws-worker/test/events/step-complete.test.ts

@@ -283,3 +283,37 @@ test('should include a timestamp', async (t) => {
   const context: any = { channel, state, onFinish: () => {} };
   await handleStepComplete(context, event);
 });
+
+test('track leaf dataclip when step has no downstream jobs', async (t) => {
+  const plan = createPlan();
+
+  const state = createRunState(plan);
+  state.activeJob = 'job-1';
+  state.activeStep = 'b';
+
+  const channel = mockChannel({
+    [STEP_COMPLETE]: () => true,
+  });
+
+  const event = { state: { x: 10 }, next: [] } as any;
+  await handleStepComplete({ channel, state } as any, event);
+
+  t.is(state.leafDataclipIds.length, 1);
+});
+
+test('do not track leaf dataclip when step has downstream jobs', async (t) => {
+  const plan = createPlan();
+
+  const state = createRunState(plan);
+  state.activeJob = 'job-1';
+  state.activeStep = 'b';
+
+  const channel = mockChannel({
+    [STEP_COMPLETE]: () => true,
+  });
+
+  const event = { state: { x: 10 }, next: ['job-2'] } as any;
+  await handleStepComplete({ channel, state } as any, event);
+
+  t.is(state.leafDataclipIds.length, 0);
+});

packages/ws-worker/test/util/create-run-state.test.ts

@@ -18,7 +18,7 @@ test('create run', (t) => {
   const run = createRunState(plan, input);
 
   t.deepEqual(run.plan, plan);
-  t.deepEqual(run.lastDataclipId, '');
+  t.deepEqual(run.leafDataclipIds, []);
   t.deepEqual(run.dataclips, {});
   t.deepEqual(run.inputDataclips, {});
   t.deepEqual(run.reasons, {});