@outputai/core 0.10.1-dev.b7b2fbe.0 → 0.10.1-next.1070949.0

package/src/hooks/index.spec.js

@@ -1,22 +1,39 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { BusEventType, ComponentType, WORKFLOW_CATALOG } from '#consts';
+import { BusEventType } from '#consts';
 
 const logErrorMock = vi.hoisted( () => vi.fn() );
 const createChildLoggerMock = vi.hoisted( () =>
   vi.fn( () => ( { error: logErrorMock } ) )
 );
-
-const onHandlers = vi.hoisted( () => ( {} ) );
-const messageBusMock = vi.hoisted( () => ( {
+const serializeErrorMock = vi.hoisted( () => vi.fn( error => ( {
+  name: error.name,
+  message: error.message,
+  stack: error.stack
+} ) ) );
+
+const mainOnHandlers = vi.hoisted( () => ( {} ) );
+const stepOnHandlers = vi.hoisted( () => ( {} ) );
+const mainEventBusMock = vi.hoisted( () => ( {
+  on: vi.fn( ( eventType, handler ) => {
+    mainOnHandlers[eventType] = handler;
+  } )
+} ) );
+const stepEventBusMock = vi.hoisted( () => ( {
+  emit: vi.fn( () => true ),
   on: vi.fn( ( eventType, handler ) => {
-    onHandlers[eventType] = handler;
+    stepOnHandlers[eventType] = handler;
   } )
 } ) );
 
 vi.mock( '#logger', () => ( { createChildLogger: createChildLoggerMock } ) );
-vi.mock( '#bus', () => ( { messageBus: messageBusMock } ) );
+vi.mock( '#helpers/error_serializer', () => ( { serializeError: serializeErrorMock } ) );
+vi.mock( '#bus', () => ( {
+  mainEventBus: mainEventBusMock,
+  stepEventBus: stepEventBusMock
+} ) );
 
 import {
+  emit,
   on,
   onActivityEnd,
   onActivityError,
@@ -27,6 +44,7 @@ import {
   onWorkflowError,
   onWorkflowStart
 } from './index.js';
+import { pendingHooks } from './pending_hooks.js';
 
 const workflowDetails = {
   workflowId: 'wf-1',
@@ -38,11 +56,6 @@ const workflowDetails = {
   attempt: 1
 };
 
-const catalogWorkflowDetails = {
-  ...workflowDetails,
-  workflowType: WORKFLOW_CATALOG
-};
-
 const activityInfo = {
   activityId: 'act-1',
   activityType: 'wf#step',
@@ -52,17 +65,15 @@ const activityInfo = {
 
 const eventDate = 1710000001234;
 
-const aggregations = {
-  cost: { total: 0 },
-  tokens: { total: 0 },
-  httpRequests: { total: 1 }
-};
-
 describe( 'hooks/index', () => {
   beforeEach( () => {
     vi.clearAllMocks();
-    Object.keys( onHandlers ).forEach( k => {
-      delete onHandlers[k];
+    pendingHooks.clear();
+    Object.keys( mainOnHandlers ).forEach( k => {
+      delete mainOnHandlers[k];
+    } );
+    Object.keys( stepOnHandlers ).forEach( k => {
+      delete stepOnHandlers[k];
     } );
   } );
 
@@ -71,9 +82,9 @@ describe( 'hooks/index', () => {
       const handler = vi.fn().mockResolvedValue( undefined );
       onError( handler );
 
-      expect( messageBusMock.on ).toHaveBeenCalledWith( BusEventType.ACTIVITY_ERROR, expect.any( Function ) );
-      expect( messageBusMock.on ).toHaveBeenCalledWith( BusEventType.WORKFLOW_ERROR, expect.any( Function ) );
-      expect( messageBusMock.on ).toHaveBeenCalledWith( BusEventType.RUNTIME_ERROR, expect.any( Function ) );
+      expect( mainEventBusMock.on ).toHaveBeenCalledWith( BusEventType.ACTIVITY_ERROR, expect.any( Function ) );
+      expect( mainEventBusMock.on ).toHaveBeenCalledWith( BusEventType.WORKFLOW_ERROR, expect.any( Function ) );
+      expect( mainEventBusMock.on ).toHaveBeenCalledWith( BusEventType.RUNTIME_ERROR, expect.any( Function ) );
     } );
 
     it( 'invokes handler with activity-shaped payload, forwarding bus fields', async () => {
@@ -81,7 +92,7 @@ describe( 'hooks/index', () => {
       onError( handler );
 
       const err = new Error( 'act-fail' );
-      await onHandlers[BusEventType.ACTIVITY_ERROR]( {
+      await mainOnHandlers[BusEventType.ACTIVITY_ERROR]( {
         eventId: 'evt-act-1',
         eventDate,
         activityInfo,
@@ -108,7 +119,7 @@ describe( 'hooks/index', () => {
       onError( handler );
 
       const err = new Error( 'wf-fail' );
-      await onHandlers[BusEventType.WORKFLOW_ERROR]( {
+      await mainOnHandlers[BusEventType.WORKFLOW_ERROR]( {
         eventId: 'evt-wf-1',
         eventDate,
         workflowDetails,
@@ -131,9 +142,17 @@ describe( 'hooks/index', () => {
       onError( handler );
 
       const error = new Error( 'rt' );
-      await onHandlers[BusEventType.RUNTIME_ERROR]( { eventId: 'evt-rt-1', eventDate, error } );
+      await mainOnHandlers[BusEventType.RUNTIME_ERROR]( { eventId: 'evt-rt-1', eventDate, error } );
 
       expect( handler ).toHaveBeenCalledWith( { eventId: 'evt-rt-1', eventDate, source: 'runtime', error } );
+      expect( logErrorMock ).toHaveBeenCalledWith( 'onError hook error', {
+        error: {
+          name: 'Error',
+          message: 'boom',
+          stack: expect.any( String )
+        }
+      } );
+      expect( serializeErrorMock ).toHaveBeenCalledWith( expect.objectContaining( { message: 'boom' } ) );
     } );
   } );
 
@@ -142,125 +161,105 @@ describe( 'hooks/index', () => {
       const handler = vi.fn().mockResolvedValue( undefined );
       onBeforeWorkerStart( handler );
 
-      expect( messageBusMock.on ).toHaveBeenCalledWith( BusEventType.WORKER_BEFORE_START, expect.any( Function ) );
-      await onHandlers[BusEventType.WORKER_BEFORE_START]();
+      expect( mainEventBusMock.on ).toHaveBeenCalledWith( BusEventType.WORKER_BEFORE_START, expect.any( Function ) );
+      await mainOnHandlers[BusEventType.WORKER_BEFORE_START]();
 
       expect( handler ).toHaveBeenCalledWith( undefined );
     } );
   } );
 
-  describe( 'onWorkflowStart', () => {
-    it( 'skips catalog workflow and forwards bus fields for real workflows', async () => {
-      const handler = vi.fn().mockResolvedValue( undefined );
-      onWorkflowStart( handler );
-
-      await Promise.resolve( onHandlers[BusEventType.WORKFLOW_START]( {
-        eventId: 'evt-ignored', eventDate, workflowDetails: catalogWorkflowDetails
-      } ) );
-      expect( handler ).not.toHaveBeenCalled();
-
-      await Promise.resolve( onHandlers[BusEventType.WORKFLOW_START]( {
-        eventId: 'evt-start-1', eventDate, workflowDetails, extra: 'passthrough'
-      } ) );
-      expect( handler ).toHaveBeenCalledWith( {
-        eventId: 'evt-start-1', eventDate, workflowDetails, extra: 'passthrough'
-      } );
+  it( 'tracks a hook callback until it settles', async () => {
+    const deferred = { resolve: null };
+    const handlerPromise = new Promise( resolve => {
+      deferred.resolve = resolve;
     } );
-  } );
+    const handler = vi.fn( () => handlerPromise );
+    onActivityEnd( handler );
 
-  describe( 'onWorkflowEnd', () => {
-    it( 'skips catalog workflow and forwards bus fields for real workflows', async () => {
-      const handler = vi.fn().mockResolvedValue( undefined );
-      onWorkflowEnd( handler );
+    const callbackPromise = mainOnHandlers[BusEventType.ACTIVITY_END]( { activityInfo } );
 
-      await Promise.resolve( onHandlers[BusEventType.WORKFLOW_END]( {
-        eventId: 'evt-ignored', eventDate, workflowDetails: catalogWorkflowDetails
-      } ) );
-      expect( handler ).not.toHaveBeenCalled();
+    expect( pendingHooks.size ).toBe( 1 );
 
-      await Promise.resolve( onHandlers[BusEventType.WORKFLOW_END]( {
-        eventId: 'evt-end-1', eventDate, workflowDetails, extra: 'passthrough'
-      } ) );
-      expect( handler ).toHaveBeenCalledWith( {
-        eventId: 'evt-end-1', eventDate, workflowDetails, extra: 'passthrough'
-      } );
-    } );
+    deferred.resolve();
+    await callbackPromise;
+    expect( pendingHooks.size ).toBe( 0 );
   } );
 
-  describe( 'onWorkflowError', () => {
-    it( 'skips catalog workflow and forwards bus fields for real workflows', async () => {
+  describe( 'workflow lifecycle hooks', () => {
+    const cases = [
+      [ 'onWorkflowStart', onWorkflowStart, BusEventType.WORKFLOW_START, {} ],
+      [ 'onWorkflowEnd', onWorkflowEnd, BusEventType.WORKFLOW_END, {} ],
+      [ 'onWorkflowError', onWorkflowError, BusEventType.WORKFLOW_ERROR, { error: new Error( 'workflow failed' ) } ]
+    ];
+
+    it.each( cases )( '%s forwards bus fields', async ( _name, registerHook, eventType, extraFields ) => {
       const handler = vi.fn().mockResolvedValue( undefined );
-      const err = new Error( 'wf' );
-      onWorkflowError( handler );
+      const payload = {
+        eventId: 'evt-workflow-1',
+        eventDate,
+        workflowDetails,
+        extra: 'passthrough',
+        ...extraFields
+      };
+      registerHook( handler );
 
-      await Promise.resolve( onHandlers[BusEventType.WORKFLOW_ERROR]( {
-        eventId: 'evt-ignored', eventDate, workflowDetails: catalogWorkflowDetails, error: err
-      } ) );
-      expect( handler ).not.toHaveBeenCalled();
+      expect( mainEventBusMock.on ).toHaveBeenCalledWith( eventType, expect.any( Function ) );
+      await mainOnHandlers[eventType]( payload );
 
-      await Promise.resolve( onHandlers[BusEventType.WORKFLOW_ERROR]( {
-        eventId: 'evt-err-1', eventDate, workflowDetails, error: err, extra: 'passthrough'
-      } ) );
-      expect( handler ).toHaveBeenCalledWith( {
-        eventId: 'evt-err-1', eventDate, workflowDetails, error: err, extra: 'passthrough'
-      } );
+      expect( handler ).toHaveBeenCalledWith( payload );
     } );
   } );
 
   describe( 'activity lifecycle hooks', () => {
     const cases = [
       [ 'onActivityStart', onActivityStart, BusEventType.ACTIVITY_START, undefined ],
-      [ 'onActivityEnd', onActivityEnd, BusEventType.ACTIVITY_END, { aggregations } ],
+      [ 'onActivityEnd', onActivityEnd, BusEventType.ACTIVITY_END ],
       [ 'onActivityError', onActivityError, BusEventType.ACTIVITY_ERROR, { error: new Error( 'activity failed' ) } ]
     ];
 
-    it.each( cases )( '%s skips internal activities and forwards bus fields', async ( _name, registerHook, eventType, extraFields = {} ) => {
+    it.each( cases )( '%s forwards internal activity bus fields', async ( _name, registerHook, eventType, extraFields = {} ) => {
       const handler = vi.fn().mockResolvedValue( undefined );
-      registerHook( handler );
-
-      expect( messageBusMock.on ).toHaveBeenCalledWith( eventType, expect.any( Function ) );
-
-      await Promise.resolve( onHandlers[eventType]( {
-        eventId: 'evt-ignored',
-        eventDate,
-        activityInfo,
-        workflowDetails,
-        outputActivityKind: ComponentType.INTERNAL_STEP,
-        ...extraFields
-      } ) );
-      expect( handler ).not.toHaveBeenCalled();
-
-      await Promise.resolve( onHandlers[eventType]( {
+      const payload = {
         eventId: 'evt-activity-1',
         eventDate,
         activityInfo,
         workflowDetails,
-        outputActivityKind: ComponentType.STEP,
+        outputActivityKind: 'internal_step',
         extra: 'passthrough',
         ...extraFields
-      } ) );
+      };
+      registerHook( handler );
 
-      expect( handler ).toHaveBeenCalledWith( {
-        eventId: 'evt-activity-1',
-        eventDate,
-        activityInfo,
-        workflowDetails,
-        outputActivityKind: ComponentType.STEP,
-        extra: 'passthrough',
-        ...extraFields
-      } );
+      expect( mainEventBusMock.on ).toHaveBeenCalledWith( eventType, expect.any( Function ) );
+      await mainOnHandlers[eventType]( payload );
+
+      expect( handler ).toHaveBeenCalledWith( payload );
     } );
   } );
 
   describe( 'on', () => {
-    it( 'subscribes to external event channel and forwards payload', async () => {
+    it( 'subscribes to SDK and user event channels and forwards payloads', async () => {
       const handler = vi.fn().mockResolvedValue( undefined );
       on( 'myEvent', handler );
 
-      expect( messageBusMock.on ).toHaveBeenCalledWith( 'external:myEvent', expect.any( Function ) );
-      await onHandlers['external:myEvent']( { foo: 1 } );
+      expect( stepEventBusMock.on ).toHaveBeenCalledWith( 'sdk:myEvent', expect.any( Function ) );
+      expect( stepEventBusMock.on ).toHaveBeenCalledWith( 'usr:myEvent', expect.any( Function ) );
+      await stepOnHandlers['sdk:myEvent']( { payload: { source: 'sdk' } } );
+      await stepOnHandlers['usr:myEvent']( { payload: { source: 'user' } } );
+
+      expect( handler ).toHaveBeenNthCalledWith( 1, { payload: { source: 'sdk' } } );
+      expect( handler ).toHaveBeenNthCalledWith( 2, { payload: { source: 'user' } } );
+    } );
+  } );
+
+  describe( 'emit', () => {
+    it( 'emits payloads on the user event channel', () => {
+      const payload = { foo: 1 };
+
+      const emitted = emit( 'myEvent', payload );
 
-      expect( handler ).toHaveBeenCalledWith( { foo: 1 } );
+      expect( stepEventBusMock.emit ).toHaveBeenCalledWith( 'usr:myEvent', payload );
+      expect( emitted ).toBe( true );
     } );
   } );
 } );

package/src/hooks/pending_hooks.js

@@ -0,0 +1,24 @@
+import { createChildLogger } from '#logger';
+
+export const pendingHooks = new Set();
+
+const log = createChildLogger( 'Hooks' );
+const flushTimeoutMs = 30_000;
+
+/**
+ * Await all pending hooks to flush for a certain time
+ */
+export const flushPendingHooks = async () => {
+  const state = { timeout: null };
+  try {
+    const flushed = await Promise.race( [
+      Promise.allSettled( [ ...pendingHooks ] ).then( _ => true ),
+      new Promise( r => state.timeout = setTimeout( () => r( false ), flushTimeoutMs ) )
+    ] );
+    if ( !flushed ) {
+      log.warn( 'Some hook callbacks exceeded the timeout and will not be awaited', { count: pendingHooks.size } );
+    }
+  } finally {
+    clearTimeout( state.timeout );
+  }
+};

package/src/hooks/pending_hooks.spec.js

@@ -0,0 +1,62 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+const logMock = vi.hoisted( () => ( { warn: vi.fn() } ) );
+
+vi.mock( '#logger', () => ( { createChildLogger: () => logMock } ) );
+
+import { flushPendingHooks, pendingHooks } from './pending_hooks.js';
+
+describe( 'pending hooks', () => {
+  beforeEach( () => {
+    pendingHooks.clear();
+    vi.clearAllMocks();
+  } );
+
+  afterEach( () => {
+    pendingHooks.clear();
+    vi.useRealTimers();
+  } );
+
+  it( 'waits for pending hooks to settle', async () => {
+    const deferred = { resolve: null };
+    const hookPromise = new Promise( resolve => {
+      deferred.resolve = resolve;
+    } );
+    pendingHooks.add( hookPromise );
+
+    const flushPromise = flushPendingHooks();
+    const state = { flushed: false };
+    flushPromise.then( () => {
+      state.flushed = true;
+    } );
+
+    await Promise.resolve();
+    expect( state.flushed ).toBe( false );
+
+    deferred.resolve();
+    await flushPromise;
+    expect( state.flushed ).toBe( true );
+  } );
+
+  it( 'stops waiting after the timeout', async () => {
+    vi.useFakeTimers();
+    pendingHooks.add( new Promise( () => {} ) );
+
+    const flushPromise = flushPendingHooks();
+    await vi.advanceTimersByTimeAsync( 30_000 );
+
+    await expect( flushPromise ).resolves.toBeUndefined();
+  } );
+
+  it( 'logs the number of hooks that exceeded the timeout', async () => {
+    vi.useFakeTimers();
+    pendingHooks.add( new Promise( () => {} ) );
+    pendingHooks.add( new Promise( () => {} ) );
+
+    const flushPromise = flushPendingHooks();
+    await vi.advanceTimersByTimeAsync( 30_000 );
+    await flushPromise;
+
+    expect( logMock.warn ).toHaveBeenCalledWith( expect.any( String ), { count: 2 } );
+  } );
+} );

package/src/interface/evaluator.d.ts

@@ -16,7 +16,7 @@ export type EvaluatorOptions = {
 /**
  * The handler function of an evaluator.
  *
- * @param input - The evaluator input; it matches the schema defined by `inputSchema`.
+ * @param input - Parsed evaluator input (`z.infer<inputSchema>`).
  *
  * @returns The result of the evaluation.
  */
@@ -30,14 +30,15 @@ export type EvaluatorFunction<
 /**
  * A wrapper around the user defined `fn` handler function.
  *
- * It has the same signature and returns the same value, calling the user function inside.
- *
- * It adds input validation based on the `inputSchema`.
+ * Callers pass values accepted by `inputSchema` (`z.input`). The wrapper parses input
+ * and invokes `fn`.
  */
-export type EvaluatorFunctionWrapper<EvaluatorFunction extends ( ...args: any ) => any> = // eslint-disable-line @typescript-eslint/no-explicit-any
-  Parameters<EvaluatorFunction> extends [infer Input] ?
-    ( input: Input ) => ReturnType<EvaluatorFunction> :
-    () => ReturnType<EvaluatorFunction>;
+export type EvaluatorFunctionWrapper<
+  InputSchema extends AnyZodSchema | undefined = undefined,
+  Result extends EvaluationResult = EvaluationResult
+> = InputSchema extends AnyZodSchema ?
+  ( input: z.input<InputSchema> ) => Promise<Result> :
+  () => Promise<Result>;
 
 /**
  * Creates an evaluation function. It is similar to a step, but must return an EvaluationResult.
@@ -67,4 +68,4 @@ export declare function evaluator<
   inputSchema: InputSchema;
   fn: EvaluatorFunction<InputSchema, Result>;
   options?: EvaluatorOptions;
-} ): EvaluatorFunctionWrapper<EvaluatorFunction<InputSchema, Result>>;
+} ): EvaluatorFunctionWrapper<InputSchema, Result>;

package/src/interface/evaluator.js

@@ -13,11 +13,6 @@ export function evaluator( { name, description, inputSchema, fn, options } ) {
     description,
     inputSchema,
     options,
-    handler: async input => {
-      validator.validateInput( input );
-      const output = await fn( input );
-      validator.validateOutput( output );
-      return output;
-    }
+    handler: async input => validator.parseOutput( await fn( validator.parseInput( input ) ) )
   } );
 }

package/src/interface/evaluator.spec.js

@@ -9,8 +9,8 @@ import {
 import { ValidationError } from '#errors';
 
 const validateDefinitionMock = vi.hoisted( () => vi.fn() );
-const validateInputMock = vi.hoisted( () => vi.fn() );
-const validateOutputMock = vi.hoisted( () => vi.fn() );
+const parseInputMock = vi.hoisted( () => vi.fn( input => input ) );
+const parseOutputMock = vi.hoisted( () => vi.fn( output => output ) );
 const validatorConstructorMock = vi.hoisted( () => vi.fn() );
 const createEvaluatorMock = vi.hoisted( () => vi.fn( ( { handler } ) => handler ) );
 
@@ -22,8 +22,8 @@ vi.mock( './validations/index.js', () => {
 
     constructor( ...args ) {
       validatorConstructorMock( ...args );
-      this.validateInput = validateInputMock;
-      this.validateOutput = validateOutputMock;
+      this.parseInput = parseInputMock;
+      this.parseOutput = parseOutputMock;
     }
   }
 
@@ -87,15 +87,15 @@ describe( 'evaluator()', () => {
 
     await expect( wrapper( { value: 'input' } ) ).resolves.toBe( output );
 
-    expect( validateInputMock ).toHaveBeenCalledWith( { value: 'input' } );
+    expect( parseInputMock ).toHaveBeenCalledWith( { value: 'input' } );
     expect( fn ).toHaveBeenCalledWith( { value: 'input' } );
-    expect( validateOutputMock ).toHaveBeenCalledWith( output );
+    expect( parseOutputMock ).toHaveBeenCalledWith( output );
   } );
 
   it( 'does not call the evaluator function when input validation throws', async () => {
     const { evaluator } = await import( './evaluator.js' );
     const error = new ValidationError( 'invalid input' );
-    validateInputMock.mockImplementationOnce( () => {
+    parseInputMock.mockImplementationOnce( () => {
       throw error;
     } );
     const fn = vi.fn();
@@ -106,13 +106,13 @@ describe( 'evaluator()', () => {
 
     await expect( wrapper( { value: 'bad' } ) ).rejects.toBe( error );
     expect( fn ).not.toHaveBeenCalled();
-    expect( validateOutputMock ).not.toHaveBeenCalled();
+    expect( parseOutputMock ).not.toHaveBeenCalled();
   } );
 
   it( 'propagates output validation errors after the evaluator function runs', async () => {
     const { evaluator } = await import( './evaluator.js' );
     const error = new ValidationError( 'invalid output' );
-    validateOutputMock.mockImplementationOnce( () => {
+    parseOutputMock.mockImplementationOnce( () => {
       throw error;
     } );
     const output = { value: 'not-an-evaluation-result' };
@@ -124,7 +124,7 @@ describe( 'evaluator()', () => {
 
     await expect( wrapper( { value: 'input' } ) ).rejects.toBe( error );
     expect( fn ).toHaveBeenCalledOnce();
-    expect( validateOutputMock ).toHaveBeenCalledWith( output );
+    expect( parseOutputMock ).toHaveBeenCalledWith( output );
   } );
 } );
 

package/src/interface/index.js

@@ -8,7 +8,7 @@ import {
 import { evaluator } from './evaluator.js';
 import { step } from './step.js';
 import { workflow } from './workflow.js';
-import { executeInParallel } from './workflow_utils.js';
+import { executeInParallel, hasErrorType } from './workflow_utils.js';
 import { sendHttpRequest, sendPostRequestAndAwaitWebhook } from './webhook.js';
 import { Logger } from './logger.js';
 
@@ -23,6 +23,7 @@ export {
   EvaluationVerdictResult,
   EvaluationFeedback,
   executeInParallel,
+  hasErrorType,
   sendHttpRequest,
   sendPostRequestAndAwaitWebhook
 };

package/src/interface/step.d.ts

@@ -15,31 +15,32 @@ export type StepOptions = {
 /**
  * The handler function of a step.
  *
- * @param input - The step input; it matches the schema defined by `inputSchema`.
+ * @param input - Parsed step input (`z.infer<inputSchema>`).
  *
- * @returns A value matching the schema defined by `outputSchema`.
+ * @returns A value accepted by `outputSchema` before parse (`z.input<outputSchema>`).
  */
 export type StepFunction<
   InputSchema extends AnyZodSchema | undefined = undefined,
   OutputSchema extends AnyZodSchema | undefined = undefined
 > = InputSchema extends AnyZodSchema ?
-  ( input: z.infer<InputSchema> ) => Promise<OutputSchema extends AnyZodSchema ? z.infer<OutputSchema> : void> :
-  () => Promise<OutputSchema extends AnyZodSchema ? z.infer<OutputSchema> : void>;
+  ( input: z.infer<InputSchema> ) => Promise<OutputSchema extends AnyZodSchema ? z.input<OutputSchema> : void> :
+  () => Promise<OutputSchema extends AnyZodSchema ? z.input<OutputSchema> : void>;
 
 /**
  * A wrapper around the user defined `fn` handler function.
  *
- * It accepts the same input and returns the same value, calling the user function inside.
- *
- * It adds input and output validation based on the `inputSchema`, `outputSchema`.
+ * Callers pass values accepted by `inputSchema` (`z.input`). The wrapper parses input,
+ * invokes `fn`, parses output, and returns `z.infer<outputSchema>`.
  *
- * @param input - The Step input; it matches the schema defined by `inputSchema`.
- * @returns A value matching the schema defined by `outputSchema`.
+ * @param input - The step input before `inputSchema` parse.
+ * @returns The step output after `outputSchema` parse.
  */
-export type StepFunctionWrapper<StepFunction extends ( ...args: any ) => any> = // eslint-disable-line @typescript-eslint/no-explicit-any
-  Parameters<StepFunction> extends [infer Input] ?
-    ( input: Input ) => ReturnType<StepFunction> :
-    () => ReturnType<StepFunction>;
+export type StepFunctionWrapper<
+  InputSchema extends AnyZodSchema | undefined = undefined,
+  OutputSchema extends AnyZodSchema | undefined = undefined
+> = InputSchema extends AnyZodSchema ?
+  ( input: z.input<InputSchema> ) => Promise<OutputSchema extends AnyZodSchema ? z.infer<OutputSchema> : void> :
+  () => Promise<OutputSchema extends AnyZodSchema ? z.infer<OutputSchema> : void>;
 
 /**
  * Creates a step.
@@ -123,7 +124,7 @@ export type StepFunctionWrapper<StepFunction extends ( ...args: any ) => any> =
  * @param params.outputSchema - Zod schema for the `fn` output
  * @param params.fn - A handler function containing the step code
  * @param params.options - Optional step options.
- * @returns The same handler function set at `fn`
+ * @returns A wrapper that parses input/output around `fn`
  */
 export declare function step<
   InputSchema extends AnyZodSchema | undefined = undefined,
@@ -135,4 +136,4 @@ export declare function step<
   outputSchema?: OutputSchema;
   fn: StepFunction<InputSchema, OutputSchema>;
   options?: StepOptions;
-} ): StepFunctionWrapper<StepFunction<InputSchema, OutputSchema>>;
+} ): StepFunctionWrapper<InputSchema, OutputSchema>;

package/src/interface/step.js

@@ -14,11 +14,6 @@ export function step( { name, description, inputSchema, outputSchema, fn, option
     inputSchema,
     outputSchema,
     options,
-    handler: async input => {
-      validator.validateInput( input );
-      const output = await fn( input );
-      validator.validateOutput( output );
-      return output;
-    }
+    handler: async input => validator.parseOutput( await fn( validator.parseInput( input ) ) )
   } );
 }