Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jumski wants to merge 1 commit into
base: 12-23-add_flowinput_to_worker_context_from_task_record
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions pkgs/client/__tests__/types/client-basic.test-d.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 () => {
Expand Down
2 changes: 1 addition & 1 deletion pkgs/client/__tests__/unit/concurrent-operations.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
91 changes: 44 additions & 47 deletions pkgs/dsl/__tests__/types/array-method.test-d.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});
});
Expand All @@ -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'
}));
});
});
Expand All @@ -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
Expand All @@ -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]);
});

Expand All @@ -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 }[];
}>();
});
Expand All @@ -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');
});
});
});
});
27 changes: 14 additions & 13 deletions pkgs/dsl/__tests__/types/context-inference.test-d.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>>();
Expand All @@ -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>>();
Expand All @@ -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>();
Expand All @@ -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
Expand Down
Loading
Loading