refactor: change step handler input to use asymmetric pattern for better type safety (#556)

# Asymmetric Step Input Pattern for Improved Type Safety and Composition

This PR introduces a significant improvement to the step input pattern across the PgFlow DSL:

- Root steps now receive flow input directly (no `run` wrapper)
- Dependent steps receive only their dependencies (flow input available via `context.flowInput`)

This change enables better functional composition where subflows can receive typed inputs without the `run` wrapper that previously blocked type matching. The asymmetric input pattern makes the API more intuitive while maintaining full type safety.

## Key Changes:

- Updated `StepInput` utility type to implement the asymmetric pattern
- Modified step handler signatures to use `flowInput` instead of `input.run`
- Added `flowInput` to the `FlowContext` interface for dependent steps
- Updated all tests to use the new pattern
- Improved type tests to verify the asymmetric behavior

This change is backward compatible at runtime but will require updating handler parameter names in client code. The improved type safety and composability are worth the minor migration effort.

Author: jumski

pkgs/client/__tests__/types/client-basic.test-d.ts

@@ -20,8 +20,8 @@ const AnalyzeWebsite = new Flow<{ url: string }>({
   baseDelay: 5,
   timeout: 10,
 })
-  .step({ slug: 'website' }, (input) => ({
-    content: `Content for ${input.run.url}`,
+  .step({ slug: 'website' }, (flowInput) => ({
+    content: `Content for ${flowInput.url}`,
   }))
   .step({ slug: 'sentiment', dependsOn: ['website'] }, (_input) => ({
     score: 0.75,
@@ -150,12 +150,16 @@ describe('PgflowClient Type Tests', () => {
       .resolves.toHaveProperty('output')
       .toEqualTypeOf<SentimentOutput | null>();
 
-    // Check input type for a step
+    // Check input type for a step - dependent steps only get their deps (no run key)
+    // flowInput is available via context.flowInput for dependent steps
     type SentimentInput = StepInput<typeof AnalyzeWebsite, 'sentiment'>;
     expectTypeOf<SentimentInput>().toMatchTypeOf<{
-      run: { url: string };
       website: { content: string };
     }>();
+    // Verify it does NOT have run key (asymmetric input)
+    expectTypeOf<SentimentInput>().not.toMatchTypeOf<{
+      run: { url: string };
+    }>();
   });
 
   it('should properly type event subscription', async () => {

pkgs/client/__tests__/unit/concurrent-operations.test.ts

@@ -23,7 +23,7 @@ import { RUN_ID, FLOW_SLUG, STEP_SLUG, startedRunSnapshot, stepStatesSample } fr
 // Create a test flow for proper typing
 const TestFlow = new Flow<{ test: string }>({ slug: 'test_flow' }).step(
   { slug: 'test_step' },
-  (input) => ({ result: input.run.test })
+  (flowInput) => ({ result: flowInput.test })
 );
 
 // Mock uuid to return predictable IDs

pkgs/dsl/__tests__/types/array-method.test-d.ts

@@ -47,49 +47,48 @@ describe('.array() method type constraints', () => {
   describe('type inference', () => {
     it('should provide correct input types for dependent steps', () => {
       new Flow<{ count: number }>({ slug: 'test' })
-        .array({ slug: 'items' }, ({ run }) => Array(run.count).fill(0).map((_, i) => i))
-        .step({ slug: 'process', dependsOn: ['items'] }, (input) => {
-          expectTypeOf(input).toMatchTypeOf<{
-            run: { count: number };
+        .array({ slug: 'items' }, (flowInput) => Array(flowInput.count).fill(0).map((_, i) => i))
+        .step({ slug: 'process', dependsOn: ['items'] }, (deps) => {
+          expectTypeOf(deps).toMatchTypeOf<{
             items: number[];
           }>();
-          return input.items.length;
+          return deps.items.length;
         });
     });
 
     it('should correctly infer element types from arrays', () => {
       new Flow<{ userId: string }>({ slug: 'test' })
         .array({ slug: 'users' }, () => [{ id: 1, name: 'John' }, { id: 2, name: 'Jane' }])
-        .step({ slug: 'count_users', dependsOn: ['users'] }, (input) => {
-          expectTypeOf(input.users).toEqualTypeOf<{ id: number; name: string }[]>();
-          expectTypeOf(input.users[0]).toMatchTypeOf<{ id: number; name: string }>();
-          return input.users.length;
+        .step({ slug: 'count_users', dependsOn: ['users'] }, (deps) => {
+          expectTypeOf(deps.users).toEqualTypeOf<{ id: number; name: string }[]>();
+          expectTypeOf(deps.users[0]).toMatchTypeOf<{ id: number; name: string }>();
+          return deps.users.length;
         });
     });
 
     it('should handle complex nested array types', () => {
       new Flow<{ depth: number }>({ slug: 'test' })
-        .array({ slug: 'matrix' }, ({ run }) => 
-          Array(run.depth).fill(0).map(() => Array(3).fill(0).map(() => ({ value: Math.random() })))
+        .array({ slug: 'matrix' }, (flowInput) =>
+          Array(flowInput.depth).fill(0).map(() => Array(3).fill(0).map(() => ({ value: Math.random() })))
         )
-        .step({ slug: 'flatten', dependsOn: ['matrix'] }, (input) => {
-          expectTypeOf(input.matrix).toEqualTypeOf<{ value: number }[][]>();
-          expectTypeOf(input.matrix[0]).toEqualTypeOf<{ value: number }[]>();
-          expectTypeOf(input.matrix[0][0]).toMatchTypeOf<{ value: number }>();
-          return input.matrix.flat();
+        .step({ slug: 'flatten', dependsOn: ['matrix'] }, (deps) => {
+          expectTypeOf(deps.matrix).toEqualTypeOf<{ value: number }[][]>();
+          expectTypeOf(deps.matrix[0]).toEqualTypeOf<{ value: number }[]>();
+          expectTypeOf(deps.matrix[0][0]).toMatchTypeOf<{ value: number }>();
+          return deps.matrix.flat();
         });
     });
 
     it('should correctly type async array handlers', () => {
       new Flow<{ url: string }>({ slug: 'test' })
-        .array({ slug: 'data' }, async ({ run }) => {
+        .array({ slug: 'data' }, async (flowInput) => {
           // Simulate async data fetching
           await new Promise(resolve => setTimeout(resolve, 1));
-          return [{ url: run.url, status: 200 }];
+          return [{ url: flowInput.url, status: 200 }];
         })
-        .step({ slug: 'validate', dependsOn: ['data'] }, (input) => {
-          expectTypeOf(input.data).toEqualTypeOf<{ url: string; status: number }[]>();
-          return input.data.every(item => item.status === 200);
+        .step({ slug: 'validate', dependsOn: ['data'] }, (deps) => {
+          expectTypeOf(deps.data).toEqualTypeOf<{ url: string; status: number }[]>();
+          return deps.data.every(item => item.status === 200);
         });
     });
   });
@@ -106,33 +105,31 @@ describe('.array() method type constraints', () => {
       new Flow<string>({ slug: 'test' })
         .array({ slug: 'items1' }, () => [1, 2, 3])
         .array({ slug: 'items2' }, () => ['a', 'b', 'c'])
-        .array({ slug: 'combined', dependsOn: ['items1'] }, (input) => {
-          expectTypeOf(input).toMatchTypeOf<{
-            run: string;
+        .array({ slug: 'combined', dependsOn: ['items1'] }, (deps) => {
+          expectTypeOf(deps).toMatchTypeOf<{
             items1: number[];
           }>();
 
           // Verify that items2 is not accessible
-          expectTypeOf(input).not.toHaveProperty('items2');
+          expectTypeOf(deps).not.toHaveProperty('items2');
 
-          return input.items1.map(String);
+          return deps.items1.map(String);
         });
     });
 
     it('should correctly type multi-dependency array steps', () => {
       new Flow<{ base: number }>({ slug: 'test' })
-        .array({ slug: 'numbers' }, ({ run }) => [run.base, run.base + 1])
+        .array({ slug: 'numbers' }, (flowInput) => [flowInput.base, flowInput.base + 1])
         .array({ slug: 'letters' }, () => ['a', 'b'])
-        .array({ slug: 'combined', dependsOn: ['numbers', 'letters'] }, (input) => {
-          expectTypeOf(input).toMatchTypeOf<{
-            run: { base: number };
+        .array({ slug: 'combined', dependsOn: ['numbers', 'letters'] }, (deps) => {
+          expectTypeOf(deps).toMatchTypeOf<{
             numbers: number[];
             letters: string[];
           }>();
-          
-          return input.numbers.map((num, i) => ({
+
+          return deps.numbers.map((num, i) => ({
             number: num,
-            letter: input.letters[i] || 'z'
+            letter: deps.letters[i] || 'z'
           }));
         });
     });
@@ -142,13 +139,13 @@ describe('.array() method type constraints', () => {
     it('should provide custom context via Flow type parameter', () => {
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       const flow = new Flow<{ id: number }, { api: { get: (id: number) => Promise<any> } }>({ slug: 'test' })
-        .array({ slug: 'fetch_data' }, (input, context) => {
+        .array({ slug: 'fetch_data' }, (flowInput, context) => {
           // No handler annotation needed! Type parameter provides context
           expectTypeOf(context.api).toEqualTypeOf<{ get: (id: number) => Promise<any> }>();
           expectTypeOf(context.env).toEqualTypeOf<Record<string, string | undefined>>();
           expectTypeOf(context.shutdownSignal).toEqualTypeOf<AbortSignal>();
 
-          return [{ id: input.run.id, data: 'mock' }];
+          return [{ id: flowInput.id, data: 'mock' }];
         });
 
       // ExtractFlowContext should include FlowContext & custom resources
@@ -164,11 +161,11 @@ describe('.array() method type constraints', () => {
 
     it('should share custom context across array and regular steps', () => {
       const flow = new Flow<{ count: number }, { generator: () => number; processor: (items: number[]) => string }>({ slug: 'test' })
-        .array({ slug: 'items' }, (input, context) => {
+        .array({ slug: 'items' }, (flowInput, context) => {
           // All steps get the same context automatically
-          return Array(input.run.count).fill(0).map(() => context.generator());
+          return Array(flowInput.count).fill(0).map(() => context.generator());
         })
-        .step({ slug: 'process' }, (input, context) => {
+        .step({ slug: 'process' }, (flowInput, context) => {
           return context.processor([1, 2, 3]);
         });
 
@@ -187,21 +184,21 @@ describe('.array() method type constraints', () => {
   describe('handler signature validation', () => {
     it('should correctly type array step handlers when using getStepDefinition', () => {
       const flow = new Flow<{ size: number }>({ slug: 'test' })
-        .array({ slug: 'data' }, (input, _context) => Array(input.run.size).fill(0).map((_, i) => ({ index: i })))
-        .step({ slug: 'dependent', dependsOn: ['data'] }, (input, _context) => input.data.length);
+        .array({ slug: 'data' }, (flowInput, _context) => Array(flowInput.size).fill(0).map((_, i) => ({ index: i })))
+        .step({ slug: 'dependent', dependsOn: ['data'] }, (deps, _context) => deps.data.length);
 
       const arrayStep = flow.getStepDefinition('data');
 
-      // Test array step handler type - handlers have 2 params (input, context)
+      // Test array step handler type - root steps receive flowInput directly (no run key)
       expectTypeOf(arrayStep.handler).toBeFunction();
-      expectTypeOf(arrayStep.handler).parameter(0).toMatchTypeOf<{ run: { size: number } }>();
+      expectTypeOf(arrayStep.handler).parameter(0).toMatchTypeOf<{ size: number }>();
       expectTypeOf(arrayStep.handler).returns.toMatchTypeOf<
         { index: number }[] | Promise<{ index: number }[]>
       >();
 
       const dependentStep = flow.getStepDefinition('dependent');
+      // Dependent steps receive deps only (no run key)
       expectTypeOf(dependentStep.handler).parameter(0).toMatchTypeOf<{
-        run: { size: number };
         data: { index: number }[];
       }>();
     });
@@ -213,20 +210,20 @@ describe('.array() method type constraints', () => {
         .array({ slug: 'items' }, () => [{ id: 1 }, { id: 2 }])
         .array({ slug: 'processed', dependsOn: ['items'] }, () => ['a', 'b']);
 
-      // Test StepInput type extraction
+      // Test StepInput type extraction - root steps get flow input directly
       type ItemsInput = StepInput<typeof flow, 'items'>;
       expectTypeOf<ItemsInput>().toMatchTypeOf<{
-        run: { userId: string };
+        userId: string;
       }>();
 
+      // Dependent steps get deps only (no run key)
       type ProcessedInput = StepInput<typeof flow, 'processed'>;
       expectTypeOf<ProcessedInput>().toMatchTypeOf<{
-        run: { userId: string };
         items: { id: number }[];
       }>();
 
       // Should not contain non-dependencies
       expectTypeOf<ProcessedInput>().not.toHaveProperty('nonExistent');
     });
   });
-});
+});

pkgs/dsl/__tests__/types/context-inference.test-d.ts

@@ -15,7 +15,7 @@ interface TestRedis {
 describe('Context Type Inference Tests', () => {
   it('should have FlowContext by default (no custom resources)', () => {
     const flow = new Flow({ slug: 'minimal_flow' })
-      .step({ slug: 'process' }, (input, context) => {
+      .step({ slug: 'process' }, (flowInput, context) => {
         // Handler automatically gets FlowContext (no annotation needed!)
         expectTypeOf(context).toMatchTypeOf<FlowContext>();
         expectTypeOf(context.env).toEqualTypeOf<Record<string, string | undefined>>();
@@ -33,7 +33,7 @@ describe('Context Type Inference Tests', () => {
 
   it('should provide custom context via Flow type parameter', () => {
     const flow = new Flow<Json, { sql: TestSql }>({ slug: 'custom_context' })
-      .step({ slug: 'query' }, (input, context) => {
+      .step({ slug: 'query' }, (flowInput, context) => {
         // No handler annotation needed! Type parameter provides context
         expectTypeOf(context.sql).toEqualTypeOf<TestSql>();
         expectTypeOf(context.env).toEqualTypeOf<Record<string, string | undefined>>();
@@ -49,13 +49,13 @@ describe('Context Type Inference Tests', () => {
 
   it('should share custom context across all steps', () => {
     const flow = new Flow<Json, { sql: TestSql; redis: TestRedis }>({ slug: 'shared_context' })
-      .step({ slug: 'query' }, (input, context) => {
+      .step({ slug: 'query' }, (flowInput, context) => {
         // All steps get the same context automatically
         expectTypeOf(context.sql).toEqualTypeOf<TestSql>();
         expectTypeOf(context.redis).toEqualTypeOf<TestRedis>();
         return { users: [] };
       })
-      .step({ slug: 'cache' }, (input, context) => {
+      .step({ slug: 'cache' }, (flowInput, context) => {
         // Second step also has access to all resources
         expectTypeOf(context.sql).toEqualTypeOf<TestSql>();
         expectTypeOf(context.redis).toEqualTypeOf<TestRedis>();
@@ -72,20 +72,21 @@ describe('Context Type Inference Tests', () => {
 
   it('should preserve existing step type inference while adding context', () => {
     const flow = new Flow<{ initial: number }, { multiplier: number }>({ slug: 'step_chain' })
-      .step({ slug: 'double' }, (input, context) => {
-        // Input inference still works
-        expectTypeOf(input.run.initial).toEqualTypeOf<number>();
+      .step({ slug: 'double' }, (flowInput, context) => {
+        // Input inference still works - root step gets flow input directly
+        expectTypeOf(flowInput.initial).toEqualTypeOf<number>();
         // Custom context available
         expectTypeOf(context.multiplier).toEqualTypeOf<number>();
-        return { doubled: input.run.initial * 2 };
+        return { doubled: flowInput.initial * 2 };
       })
-      .step({ slug: 'format', dependsOn: ['double'] }, (input, context) => {
-        // Dependent step has access to previous step output
-        expectTypeOf(input.run.initial).toEqualTypeOf<number>();
-        expectTypeOf(input.double.doubled).toEqualTypeOf<number>();
+      .step({ slug: 'format', dependsOn: ['double'] }, (deps, context) => {
+        // Dependent step has access to previous step output via deps
+        expectTypeOf(deps.double.doubled).toEqualTypeOf<number>();
         // And still has custom context
         expectTypeOf(context.multiplier).toEqualTypeOf<number>();
-        return { formatted: String(input.double.doubled) };
+        // Access flow input via context.flowInput
+        expectTypeOf(context.flowInput.initial).toEqualTypeOf<number>();
+        return { formatted: String(deps.double.doubled) };
       });
 
     // Context includes custom resources