feat: add flow_input to StepTaskRecord and context objects (#555)

# Add flow_input to StepTaskRecord and expose it in worker context

This PR adds the original flow input to the `StepTaskRecord` type and exposes it throughout the worker context. This allows dependent steps to access the original flow input parameters, which is useful for maintaining context across multi-step flows.

The changes:
- Add `flow_input` field to `StepTaskRecord` in core types
- Expose the flow input in the `StepTaskExecution` and `StepTaskWithMessage` interfaces
- Pass the flow input through the task poller to the worker context
- Update test helpers to include flow input in context creation
- Add tests to verify flow input is properly passed through the context

This enhancement enables step implementations to reference the original parameters that started the flow, improving the ability to build steps that depend on initial configuration.

Author: jumski

pkgs/core/src/types.ts

@@ -25,6 +25,7 @@ export type StepTaskRecord<TFlow extends AnyFlow> = {
     task_index: number;
     input: Simplify<StepInput<TFlow, StepSlug>>;
     msg_id: number;
+    flow_input: ExtractFlowInput<TFlow>;
   };
 }[Extract<keyof ExtractFlowSteps<TFlow>, string>];
 

pkgs/edge-worker/deno.lock

@@ -1,5 +1,5 @@
 /* DSL‐level ------------------------------------------------------------ */
-import type { BaseContext, AnyFlow, AllStepInputs } from '@pgflow/dsl';
+import type { BaseContext, AnyFlow, AllStepInputs, ExtractFlowInput } from '@pgflow/dsl';
 import type { Json } from './types.js';
 import type { PgmqMessageRecord } from '../queue/types.js';
 import type { StepTaskRecord } from '../flow/types.js';
@@ -30,6 +30,7 @@ export interface StepTaskExecution<TFlow extends AnyFlow = AnyFlow> {
   rawMessage: PgmqMessageRecord<AllStepInputs<TFlow>>;
   stepTask  : StepTaskRecord<TFlow>;
   workerConfig: Readonly<Omit<FlowWorkerConfig, 'sql'>>;
+  flowInput: ExtractFlowInput<TFlow>;
 }
 
 /** Message handler context for any platform */
@@ -52,6 +53,7 @@ export interface StepTaskWithMessage<TFlow extends AnyFlow> {
   msg_id : number;
   message: PgmqMessageRecord<AllStepInputs<TFlow>>;
   task   : StepTaskRecord<TFlow>;
+  flowInput: ExtractFlowInput<TFlow>;
 }
 
 import { deepClone, deepFreeze } from './deepUtils.js';

pkgs/edge-worker/src/core/context.ts

@@ -103,7 +103,8 @@ export class StepTaskPoller<TFlow extends AnyFlow>
           return {
             message,
             task,
-            msg_id: task.msg_id
+            msg_id: task.msg_id,
+            flowInput: task.flow_input
           };
         })
         .filter((item): item is StepTaskWithMessage<TFlow> => item !== null);

pkgs/edge-worker/src/flow/StepTaskPoller.ts

@@ -132,6 +132,7 @@ export function createFlowWorker<TFlow extends AnyFlow, TResources extends Recor
       rawMessage: taskWithMessage.message,
       stepTask: taskWithMessage.task,
       workerConfig: frozenWorkerConfig, // Reuse cached frozen config
+      flowInput: taskWithMessage.flowInput, // Original flow input for dependent steps
 
       // Platform-specific resources (generic)
       ...platformAdapter.platformResources

pkgs/edge-worker/src/flow/createFlowWorker.ts

@@ -1,6 +1,6 @@
 import type postgres from 'postgres';
 import type { SupabaseClient } from '@supabase/supabase-js';
-import type { AnyFlow, AllStepInputs } from '@pgflow/dsl';
+import type { AnyFlow, AllStepInputs, ExtractFlowInput } from '@pgflow/dsl';
 import type { SupabaseResources } from '@pgflow/dsl/supabase';
 import type { Json } from '../core/types.js';
 import type { PgmqMessageRecord } from '../queue/types.js';
@@ -83,9 +83,10 @@ export function createSupabaseStepTaskContext<TFlow extends AnyFlow>(params: {
   abortSignal: AbortSignal;
   stepTask: StepTaskRecord<TFlow>;
   rawMessage: PgmqMessageRecord<AllStepInputs<TFlow>>;
+  flowInput: ExtractFlowInput<TFlow>;
   workerConfig?: Readonly<Omit<FlowWorkerConfig, 'sql'>>;
 }): SupabaseStepTaskContext<TFlow> {
-  const { env, sql, abortSignal, stepTask, rawMessage, workerConfig } = params;
+  const { env, sql, abortSignal, stepTask, rawMessage, flowInput, workerConfig } = params;
 
   // Validate required environment variables
   const supabaseUrl = env.SUPABASE_URL;
@@ -126,6 +127,7 @@ export function createSupabaseStepTaskContext<TFlow extends AnyFlow>(params: {
     rawMessage,
     stepTask,
     workerConfig: defaultWorkerConfig,
+    flowInput,
 
     // Supabase-specific resources (always present)
     sql,

pkgs/edge-worker/src/test/test-helpers.ts

@@ -58,6 +58,7 @@ Deno.test(
       step_slug: 'test_step',
       task_index: 0,
       input: { run: { data: 'test data' } },
+      flow_input: { data: 'test data' },
     };
 
     // Create context with mock task and message using proper flow worker context creation
@@ -77,6 +78,7 @@ Deno.test(
         msg_id: 123,
         message: mockMessage,
         task: mockTask,
+        flowInput: { data: 'test data' },
       },
     });
 
@@ -120,6 +122,7 @@ Deno.test(
       step_slug: 'legacy_step',
       task_index: 0,
       input: { run: { value: 42 } },
+      flow_input: { value: 42 },
     };
 
     // Get the step handler
@@ -142,6 +145,7 @@ Deno.test(
         msg_id: 456,
         message: mockMessage,
         task: mockTask,
+        flowInput: { value: 42 },
       },
     });
 
@@ -187,13 +191,15 @@ Deno.test(
       step_slug: 'check_raw',
       task_index: 0,
       input: { run: {} },
+      flow_input: {},
     };
 
     // Create context - for this test we need a mock taskWithMessage
     const mockTaskWithMessage = {
       msg_id: 789,
       message: mockMessage,
       task: mockTask,
+      flowInput: {},
     };
 
     const context = createFlowWorkerContext({
@@ -247,13 +253,15 @@ Deno.test(
       step_slug: 'check_clients',
       task_index: 0,
       input: { run: {} },
+      flow_input: {},
     };
 
     // Create context with Supabase env vars
     const mockTaskWithMessage = {
       msg_id: 999,
       message: mockMessage,
       task: mockTask,
+      flowInput: {},
     };
 
     const context = createFlowWorkerContext({
@@ -328,6 +336,7 @@ Deno.test(
       step_slug: 'fetch_data',
       task_index: 0,
       input: { run: { id: 123 } },
+      flow_input: { id: 123 },
     };
 
     const context = createFlowWorkerContext({
@@ -338,6 +347,7 @@ Deno.test(
         msg_id: 456,
         message: mockMessageForComplex,
         task: mockTaskForComplex,
+        flowInput: { id: 123 },
       },
     });
 

pkgs/edge-worker/tests/integration/stepTaskExecutorContext.test.ts

@@ -47,13 +47,18 @@ const mockStepMessage: PgmqMessageRecord<{ run: { test: string } }> = {
   headers: null,
 };
 
+// Mock flow input
+const mockFlowInput = { test: 'flow-input' };
+
 // Mock step task (using generic typing)
 const mockStepTask = {
   flow_slug: 'test-flow',
   run_id: 'run-456',
   step_slug: 'test-step',
   input: { run: { test: 'input' } },
-  msg_id: 123
+  msg_id: 123,
+  flow_input: mockFlowInput,
+  task_index: 0
 } as StepTaskRecord<never>;
 
 Deno.test('createSupabaseMessageContext - creates context with all Supabase resources', () => {
@@ -104,15 +109,17 @@ Deno.test('createSupabaseStepTaskContext - creates context with step task', () =
     abortSignal: mockAbortSignal,
     stepTask: mockStepTask,
     rawMessage: mockStepMessage,
+    flowInput: mockFlowInput,
   });
-  
+
   // Check all properties exist
   assertEquals(context.env, fullEnv);
   assertEquals(context.sql, mockSql);
   assertEquals(context.shutdownSignal, mockAbortSignal);
   assertEquals(context.stepTask, mockStepTask);
   assertEquals(context.rawMessage, mockStepMessage);
-  
+  assertEquals(context.flowInput, mockFlowInput);
+
   // Supabase client should always be present
   assertExists(context.supabase);
 });

pkgs/edge-worker/tests/unit/contextUtils.test.ts

@@ -41,6 +41,7 @@ export const stepTaskRecord: StepTaskRecord<typeof ExampleFlow> = {
     // normalStep: { doubledValueArray: [1, 2, 3] }, --- this should be an error
   },
   msg_id: 1,
+  flow_input: { value: 23 },
 };
 
 // export const yolo: { value: number } = { value: 23, otherValue: 'yolo' };