From b0de0cbba74a4a7e5f11f2fa52211a85cdaba45e Mon Sep 17 00:00:00 2001
From: unnoq <contact@unnoq.com>
Date: Sat, 5 Apr 2025 15:54:14 +0700
Subject: [PATCH] fix(server): onSuccess, onError, ... cannot use as a
 middleware

---
 packages/shared/src/interceptor.test-d.ts | 36 +++++++++++++++++++++++
 packages/shared/src/interceptor.ts        | 18 ++++++------
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/packages/shared/src/interceptor.test-d.ts b/packages/shared/src/interceptor.test-d.ts
index 73afd6dbd..e635048a5 100644
--- a/packages/shared/src/interceptor.test-d.ts
+++ b/packages/shared/src/interceptor.test-d.ts
@@ -1,5 +1,8 @@
+import type { Context } from '../../server/src/context'
+import type { MiddlewareNextFn, MiddlewareResult } from '../../server/src/middleware'
 import type { Interceptor } from './interceptor'
 import type { PromiseWithError } from './types'
+import { os } from '../../server/src/builder'
 import { onError, onFinish, onStart, onSuccess } from './interceptor'
 
 it('onStart', () => {
@@ -8,6 +11,13 @@ it('onStart', () => {
     expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
     expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
   })
+
+  os.$context<{ something: string }>().use(onStart(({ context, next }) => {
+    expectTypeOf(context).toEqualTypeOf<{ something: string }>()
+    expectTypeOf(next).toEqualTypeOf<MiddlewareNextFn<unknown>>()
+  })).handler(({ context }) => {
+    expectTypeOf(context).toMatchTypeOf<{ something: string }>()
+  })
 })
 
 it('onSuccess', () => {
@@ -18,6 +28,14 @@ it('onSuccess', () => {
     expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
     expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
   })
+
+  os.$context<{ something: string }>().use(onSuccess((data, { context, next }) => {
+    expectTypeOf(data).toEqualTypeOf <Awaited<MiddlewareResult<Context, unknown>>>()
+    expectTypeOf(context).toEqualTypeOf<{ something: string }>()
+    expectTypeOf(next).toEqualTypeOf<MiddlewareNextFn<unknown>>()
+  })).handler(({ context }) => {
+    expectTypeOf(context).toMatchTypeOf<{ something: string }>()
+  })
 })
 
 it('onError', () => {
@@ -28,6 +46,14 @@ it('onError', () => {
     expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
     expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
   })
+
+  os.$context<{ something: string }>().use(onError((error, { context, next }) => {
+    expectTypeOf(error).toEqualTypeOf<Error>()
+    expectTypeOf(context).toEqualTypeOf<{ something: string }>()
+    expectTypeOf(next).toEqualTypeOf<MiddlewareNextFn<unknown>>()
+  })).handler(({ context }) => {
+    expectTypeOf(context).toMatchTypeOf<{ something: string }>()
+  })
 })
 
 it('onFinish', () => {
@@ -38,4 +64,14 @@ it('onFinish', () => {
     expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
     expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
   })
+
+  os.$context<{ something: string }>().use(onFinish(() => { }))
+
+  os.$context<{ something: string }>().use(onFinish((state, { context, next }) => {
+    expectTypeOf(state).toEqualTypeOf<[Awaited<MiddlewareResult<Context, unknown>>, null, 'success'] | [undefined, Error, 'error']>()
+    expectTypeOf(context).toEqualTypeOf<{ something: string }>()
+    expectTypeOf(next).toEqualTypeOf<MiddlewareNextFn<unknown>>()
+  })).handler(({ context }) => {
+    expectTypeOf(context).toMatchTypeOf<{ something: string }>()
+  })
 })
diff --git a/packages/shared/src/interceptor.ts b/packages/shared/src/interceptor.ts
index 26ed5a80c..cd1a845d1 100644
--- a/packages/shared/src/interceptor.ts
+++ b/packages/shared/src/interceptor.ts
@@ -15,14 +15,14 @@ export type Interceptor<
   TOptions extends InterceptableOptions,
   TResult,
   TError,
-> = (options: InterceptorOptions<TOptions, TResult, TError>) => PromiseWithError<TResult, TError>
+> = (options: InterceptorOptions<TOptions, TResult, TError>) => Promisable<TResult>
 
 /**
  * Can used for interceptors or middlewares
  */
-export function onStart<TOptions extends { next(): any }, TRest extends any[]>(
+export function onStart<T, TOptions extends { next(): any }, TRest extends any[]>(
   callback: NoInfer<(options: TOptions, ...rest: TRest) => Promisable<void>>,
-): (options: TOptions, ...rest: TRest) => Promise<Awaited<ReturnType<TOptions['next']>>> {
+): (options: TOptions, ...rest: TRest) => T | Promise<Awaited<ReturnType<TOptions['next']>>> {
   return async (options, ...rest) => {
     await callback(options, ...rest)
     return await options.next()
@@ -32,9 +32,9 @@ export function onStart<TOptions extends { next(): any }, TRest extends any[]>(
 /**
  * Can used for interceptors or middlewares
  */
-export function onSuccess<TOptions extends { next(): any }, TRest extends any[]>(
+export function onSuccess<T, TOptions extends { next(): any }, TRest extends any[]>(
   callback: NoInfer<(result: Awaited<ReturnType<TOptions['next']>>, options: TOptions, ...rest: TRest) => Promisable<void>>,
-): (options: TOptions, ...rest: TRest) => Promise<Awaited<ReturnType<TOptions['next']>>> {
+): (options: TOptions, ...rest: TRest) => T | Promise<Awaited<ReturnType<TOptions['next']>>> {
   return async (options, ...rest) => {
     const result = await options.next()
     await callback(result, options, ...rest)
@@ -45,13 +45,13 @@ export function onSuccess<TOptions extends { next(): any }, TRest extends any[]>
 /**
  * Can used for interceptors or middlewares
  */
-export function onError<TOptions extends { next(): any }, TRest extends any[]>(
+export function onError<T, TOptions extends { next(): any }, TRest extends any[]>(
   callback: NoInfer<(
     error: ReturnType<TOptions['next']> extends PromiseWithError<any, infer E> ? E : ThrowableError,
     options: TOptions,
     ...rest: TRest
   ) => Promisable<void>>,
-): (options: TOptions, ...rest: TRest) => Promise<Awaited<ReturnType<TOptions['next']>>> {
+): (options: TOptions, ...rest: TRest) => T | Promise<Awaited<ReturnType<TOptions['next']>>> {
   return async (options, ...rest) => {
     try {
       return await options.next()
@@ -68,7 +68,7 @@ export type OnFinishState<TResult, TError> = [TResult, null, 'success'] | [undef
 /**
  * Can used for interceptors or middlewares
  */
-export function onFinish<TOptions extends { next(): any }, TRest extends any[]>(
+export function onFinish<T, TOptions extends { next(): any }, TRest extends any[]>(
   callback: NoInfer<(
     state: OnFinishState<
       Awaited<ReturnType<TOptions['next']>>,
@@ -77,7 +77,7 @@ export function onFinish<TOptions extends { next(): any }, TRest extends any[]>(
     options: TOptions,
     ...rest: TRest
   ) => Promisable<void>>,
-): (options: TOptions, ...rest: TRest) => Promise<Awaited<ReturnType<TOptions['next']>>> {
+): (options: TOptions, ...rest: TRest) => T | Promise<Awaited<ReturnType<TOptions['next']>>> {
   let state: any
 
   return async (options, ...rest) => {