Enforce JSON-safe step outputs without narrowing inference

Author: jumski

.changeset/enforce-json-step-outputs.md

@@ -0,0 +1,5 @@
+---
+'@pgflow/dsl': patch
+---
+
+Enforce JSON-compatible step outputs at .step() construction time

pkgs/dsl/__tests__/types/step-method-json-output.test-d.ts

@@ -0,0 +1,43 @@
+import { Flow, type StepOutput } from '../../src/index.js';
+import { describe, expectTypeOf, it } from 'vitest';
+
+describe('.step() JSON output constraints', () => {
+  it('accepts JSON-compatible step outputs', () => {
+    new Flow<{ x: string }>({ slug: 'test' })
+      .step({ slug: 'valid1' }, () => ({ ok: true }))
+      .step({ slug: 'valid2' }, () => ({ maybe: null as string | null }))
+      .step({ slug: 'valid3' }, () => ['a', 'b', 'c'])
+      .step({ slug: 'valid4' }, () => null)
+      .step({ slug: 'valid5' }, () => 123);
+  });
+
+  it('preserves widened output inference for step outputs', () => {
+    const flow = new Flow<{ x: string }>({ slug: 'test_widened' }).step(
+      { slug: 'numberStep' },
+      () => 123
+    );
+
+    expectTypeOf<
+      StepOutput<typeof flow, 'numberStep'>
+    >().toEqualTypeOf<number>();
+  });
+
+  it('rejects undefined in step outputs', () => {
+    new Flow<{ x: string }>({ slug: 'test_invalid' })
+      // @ts-expect-error - undefined is not Json
+      .step({ slug: 'invalid1' }, () => undefined)
+      // @ts-expect-error - property type includes undefined (not Json)
+      .step({ slug: 'invalid2' }, () => ({
+        ok: true,
+        maybe: undefined as string | undefined,
+      }));
+  });
+
+  it('rejects other non-JSON output types', () => {
+    new Flow<{ x: string }>({ slug: 'test_invalid_other' })
+      // @ts-expect-error - symbol is not Json
+      .step({ slug: 'invalid3' }, () => Symbol('x'))
+      // @ts-expect-error - function is not Json
+      .step({ slug: 'invalid4' }, () => () => 'nope');
+  });
+});

pkgs/dsl/src/dsl.ts

@@ -40,6 +40,20 @@ type AwaitedReturn<T> = T extends (...args: any[]) => Promise<infer R>
   ? R
   : never;
 
+type WidenJson<T> = T extends string
+  ? string
+  : T extends number
+  ? number
+  : T extends boolean
+  ? boolean
+  : T extends null
+  ? null
+  : T extends readonly (infer U)[]
+  ? WidenJson<U>[]
+  : T extends object
+  ? { [K in keyof T]: WidenJson<T[K]> }
+  : T;
+
 // ========================
 // ENVIRONMENT TYPE SYSTEM
 // ========================
@@ -716,7 +730,10 @@ export class Flow<
   // Overload 1: Root step without conditions
   step<
     Slug extends string,
-    TOutput,
+    THandler extends (
+      flowInput: TFlowInput,
+      context: FlowContext<TEnv, TFlowInput> & TContext
+    ) => Json | Promise<Json>,
     TRetries extends WhenExhaustedMode | undefined = undefined
   >(
     opts: Simplify<
@@ -727,16 +744,13 @@ export class Flow<
       } & WithoutCondition &
         Omit<BaseStepRuntimeOptions, 'whenExhausted'>
     >,
-    handler: (
-      flowInput: TFlowInput,
-      context: FlowContext<TEnv, TFlowInput> & TContext
-    ) => TOutput | Promise<TOutput>
+    handler: THandler
   ): Flow<
     TFlowInput,
     TContext,
     Steps & {
       [K in Slug]: StepMeta<
-        Awaited<TOutput>,
+        WidenJson<AwaitedReturn<THandler>>,
         TRetries extends 'skip' | 'skip-cascade' ? TRetries : false
       >;
     },
@@ -747,7 +761,10 @@ export class Flow<
   // Overload 2: Root step with condition and omitted whenUnmet defaults to 'skip'
   step<
     Slug extends string,
-    TOutput,
+    THandler extends (
+      flowInput: TFlowInput,
+      context: FlowContext<TEnv, TFlowInput> & TContext
+    ) => Json | Promise<Json>,
     TRetries extends WhenExhaustedMode | undefined = undefined
   >(
     opts: Simplify<
@@ -761,15 +778,12 @@ export class Flow<
       ) &
         Omit<BaseStepRuntimeOptions, 'whenExhausted'>
     >,
-    handler: (
-      flowInput: TFlowInput,
-      context: FlowContext<TEnv, TFlowInput> & TContext
-    ) => TOutput | Promise<TOutput>
+    handler: THandler
   ): Flow<
     TFlowInput,
     TContext,
     Steps & {
-      [K in Slug]: StepMeta<Awaited<TOutput>, 'skip'>;
+      [K in Slug]: StepMeta<WidenJson<AwaitedReturn<THandler>>, 'skip'>;
     },
     StepDependencies & { [K in Slug]: [] },
     TEnv
@@ -778,7 +792,10 @@ export class Flow<
   // Overload 3: Root step with explicit whenUnmet
   step<
     Slug extends string,
-    TOutput,
+    THandler extends (
+      flowInput: TFlowInput,
+      context: FlowContext<TEnv, TFlowInput> & TContext
+    ) => Json | Promise<Json>,
     TWhenUnmet extends WhenUnmetMode,
     TRetries extends WhenExhaustedMode | undefined = undefined
   >(
@@ -793,16 +810,13 @@ export class Flow<
       ) &
         Omit<BaseStepRuntimeOptions, 'whenExhausted'>
     >,
-    handler: (
-      flowInput: TFlowInput,
-      context: FlowContext<TEnv, TFlowInput> & TContext
-    ) => TOutput | Promise<TOutput>
+    handler: THandler
   ): Flow<
     TFlowInput,
     TContext,
     Steps & {
       [K in Slug]: StepMeta<
-        Awaited<TOutput>,
+        WidenJson<AwaitedReturn<THandler>>,
         TWhenUnmet extends 'skip' | 'skip-cascade'
           ? TWhenUnmet
           : TRetries extends 'skip' | 'skip-cascade'
@@ -818,7 +832,10 @@ export class Flow<
   step<
     Slug extends string,
     Deps extends Extract<keyof Steps, string>,
-    TOutput,
+    THandler extends (
+      deps: Simplify<DepsWithOptionalSkippable<Steps, Deps>>,
+      context: FlowContext<TEnv, TFlowInput> & TContext
+    ) => Json | Promise<Json>,
     TRetries extends WhenExhaustedMode | undefined = undefined
   >(
     opts: Simplify<
@@ -829,16 +846,13 @@ export class Flow<
       } & WithoutCondition &
         Omit<BaseStepRuntimeOptions, 'whenExhausted'>
     >,
-    handler: (
-      deps: Simplify<DepsWithOptionalSkippable<Steps, Deps>>,
-      context: FlowContext<TEnv, TFlowInput> & TContext
-    ) => TOutput | Promise<TOutput>
+    handler: THandler
   ): Flow<
     TFlowInput,
     TContext,
     Steps & {
       [K in Slug]: StepMeta<
-        Awaited<TOutput>,
+        WidenJson<AwaitedReturn<THandler>>,
         TRetries extends 'skip' | 'skip-cascade' ? TRetries : false
       >;
     },
@@ -850,7 +864,10 @@ export class Flow<
   step<
     Slug extends string,
     Deps extends Extract<keyof Steps, string>,
-    TOutput,
+    THandler extends (
+      deps: Simplify<DepsWithOptionalSkippable<Steps, Deps>>,
+      context: FlowContext<TEnv, TFlowInput> & TContext
+    ) => Json | Promise<Json>,
     TRetries extends WhenExhaustedMode | undefined = undefined
   >(
     opts: Simplify<
@@ -870,15 +887,12 @@ export class Flow<
       ) &
         Omit<BaseStepRuntimeOptions, 'whenExhausted'>
     >,
-    handler: (
-      deps: Simplify<DepsWithOptionalSkippable<Steps, Deps>>,
-      context: FlowContext<TEnv, TFlowInput> & TContext
-    ) => TOutput | Promise<TOutput>
+    handler: THandler
   ): Flow<
     TFlowInput,
     TContext,
     Steps & {
-      [K in Slug]: StepMeta<Awaited<TOutput>, 'skip'>;
+      [K in Slug]: StepMeta<WidenJson<AwaitedReturn<THandler>>, 'skip'>;
     },
     StepDependencies & { [K in Slug]: Deps[] },
     TEnv
@@ -888,7 +902,10 @@ export class Flow<
   step<
     Slug extends string,
     Deps extends Extract<keyof Steps, string>,
-    TOutput,
+    THandler extends (
+      deps: Simplify<DepsWithOptionalSkippable<Steps, Deps>>,
+      context: FlowContext<TEnv, TFlowInput> & TContext
+    ) => Json | Promise<Json>,
     TWhenUnmet extends WhenUnmetMode,
     TRetries extends WhenExhaustedMode | undefined = undefined
   >(
@@ -907,16 +924,13 @@ export class Flow<
       ) &
         Omit<BaseStepRuntimeOptions, 'whenExhausted'>
     >,
-    handler: (
-      deps: Simplify<DepsWithOptionalSkippable<Steps, Deps>>,
-      context: FlowContext<TEnv, TFlowInput> & TContext
-    ) => TOutput | Promise<TOutput>
+    handler: THandler
   ): Flow<
     TFlowInput,
     TContext,
     Steps & {
       [K in Slug]: StepMeta<
-        Awaited<TOutput>,
+        WidenJson<AwaitedReturn<THandler>>,
         TWhenUnmet extends 'skip' | 'skip-cascade'
           ? TWhenUnmet
           : TRetries extends 'skip' | 'skip-cascade'