fix: .next broken when calling multiple times (#337)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Refactor**
- Streamlined internal processing logic to ensure more consistent and
reliable handling of sequential operations.
- **Tests**
- Expanded test coverage by adding scenarios that validate the system's
behavior when executing multiple sequential steps.
- Improved test setup for more accurate state management between test
runs.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: dinwwwh

packages/server/src/procedure-client.test.ts

@@ -45,7 +45,7 @@ const procedureCases = [
 ] as const
 
 beforeEach(() => {
-  vi.clearAllMocks()
+  vi.resetAllMocks()
 })
 
 describe.each(procedureCases)('createProcedureClient - case %s', async (_, procedure) => {
@@ -479,6 +479,37 @@ describe.each(procedureCases)('createProcedureClient - case %s', async (_, proce
     expect(context).toBeCalledTimes(1)
     expect(context).toBeCalledWith({ cache: true })
   })
+
+  it('can multiple .next calls', async () => {
+    const client = createProcedureClient(procedure)
+
+    preMid1.mockImplementationOnce(async ({ next }, input, output) => output([(await next()).output, (await next()).output, (await next()).output, (await next()).output]))
+
+    let index = 0
+
+    preMid2.mockImplementation(({ next }) => next({ context: { preMid2: index++ } }))
+    postMid1.mockImplementation(({ next }) => next({ context: { postMid1: index++ } }))
+
+    await expect(client({ val: '123' })).resolves.toEqual([{ val: 123 }, { val: 123 }, { val: 123 }, { val: 123 }])
+
+    expect(preMid1).toBeCalledTimes(1)
+    expect(preMid2).toBeCalledTimes(4)
+    expect(postMid1).toBeCalledTimes(4)
+    expect(postMid2).toBeCalledTimes(4)
+    expect(handler).toBeCalledTimes(4)
+
+    expect((handler as any).mock.calls[0][0].context.preMid2).toBe(0)
+    expect((handler as any).mock.calls[0][0].context.postMid1).toBe(1)
+
+    expect((handler as any).mock.calls[1][0].context.preMid2).toBe(2)
+    expect((handler as any).mock.calls[1][0].context.postMid1).toBe(3)
+
+    expect((handler as any).mock.calls[2][0].context.preMid2).toBe(4)
+    expect((handler as any).mock.calls[2][0].context.postMid1).toBe(5)
+
+    expect((handler as any).mock.calls[3][0].context.preMid2).toBe(6)
+    expect((handler as any).mock.calls[3][0].context.postMid1).toBe(7)
+  })
 })
 
 it('still work without InputSchema', async () => {

packages/server/src/procedure-client.ts

@@ -1,7 +1,6 @@
 import type { Client, ClientContext } from '@orpc/client'
 import type { Interceptor, MaybeOptionalOptions, Value } from '@orpc/shared'
 import type { Lazyable } from './lazy'
-import type { MiddlewareNextFn } from './middleware'
 import type { AnyProcedure, Procedure, ProcedureHandlerOptions } from './procedure'
 import { ORPCError } from '@orpc/client'
 import { type AnySchema, type ErrorFromErrorMap, type ErrorMap, type InferSchemaInput, type InferSchemaOutput, type Meta, ValidationError } from '@orpc/contract'
@@ -162,44 +161,37 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced
   const middlewares = procedure['~orpc'].middlewares
   const inputValidationIndex = Math.min(Math.max(0, procedure['~orpc'].inputValidationIndex), middlewares.length)
   const outputValidationIndex = Math.min(Math.max(0, procedure['~orpc'].outputValidationIndex), middlewares.length)
-  let currentIndex = 0
-  let currentContext = options.context
-  let currentInput = options.input
 
-  const next: MiddlewareNextFn<any> = async (...[nextOptions]) => {
-    const index = currentIndex
-    const midContext = nextOptions?.context ?? {} as any
-
-    currentIndex += 1
-    currentContext = mergeCurrentContext(currentContext, midContext)
+  const next = async (index: number, context: Context, input: unknown): Promise<unknown> => {
+    let currentInput = input
 
     if (index === inputValidationIndex) {
       currentInput = await validateInput(procedure, currentInput)
     }
 
     const mid = middlewares[index]
 
-    const result = mid
-      ? {
-          context: midContext,
-          output: (await mid({ ...options, context: currentContext, next }, currentInput, middlewareOutputFn)).output,
-        }
-      : {
-          context: midContext,
-          output: await procedure['~orpc'].handler({ ...options, context: currentContext, input: currentInput }),
-        }
+    const output = mid
+      ? (await mid({
+          ...options,
+          context,
+          next: async (...[nextOptions]) => {
+            const nextContext: Context = nextOptions?.context ?? {}
 
-    if (index === outputValidationIndex) {
-      const validatedOutput = await validateOutput(procedure, result.output)
+            return {
+              output: await next(index + 1, mergeCurrentContext(context, nextContext), currentInput),
+              context: nextContext,
+            }
+          },
+        }, currentInput, middlewareOutputFn)).output
+      : await procedure['~orpc'].handler({ ...options, context, input: currentInput })
 
-      return {
-        context: result.context,
-        output: validatedOutput,
-      }
+    if (index === outputValidationIndex) {
+      return await validateOutput(procedure, output)
     }
 
-    return result
+    return output
   }
 
-  return (await next()).output
+  return next(0, options.context, options.input)
 }

packages/shared/src/interceptor.test.ts

@@ -115,6 +115,41 @@ describe('intercept', () => {
       next: 'hello2',
     })
   })
+
+  it('can multiple .next calls', async () => {
+    interceptor2.mockReturnValueOnce(Promise.resolve('__interceptor2__'))
+
+    const result = await intercept(
+      [
+        async ({ next }) => [await next(), await next(), await next()],
+        interceptor1,
+        interceptor2,
+      ],
+      {
+        foo: 'bar',
+      },
+      main,
+    )
+
+    expect(result).toEqual(['__interceptor2__', '__main__', '__main__'])
+
+    expect(interceptor1).toHaveBeenCalledTimes(3)
+    expect(interceptor1).toHaveBeenCalledWith({
+      foo: 'bar',
+      next: expect.any(Function),
+    })
+
+    expect(interceptor2).toHaveBeenCalledTimes(3)
+    expect(interceptor2).toHaveBeenCalledWith({
+      foo: 'bar',
+      next: expect.any(Function),
+    })
+
+    expect(main).toHaveBeenCalledTimes(2)
+    expect(main).toHaveBeenCalledWith({
+      foo: 'bar',
+    })
+  })
 })
 
 describe('onStart / onSuccess / onError / onFinish', () => {

packages/shared/src/interceptor.ts

@@ -103,20 +103,18 @@ export async function intercept<TOptions extends InterceptableOptions, TResult,
   options: NoInfer<TOptions>,
   main: NoInfer<(options: TOptions) => Promisable<TResult>>,
 ): Promise<TResult> {
-  let index = 0
-
-  const next = async (options: TOptions): Promise<TResult> => {
-    const interceptor = interceptors[index++]
+  const next = async (options: TOptions, index: number): Promise<TResult> => {
+    const interceptor = interceptors[index]
 
     if (!interceptor) {
       return await main(options)
     }
 
     return await interceptor({
       ...options,
-      next: (newOptions: TOptions = options) => next(newOptions),
+      next: (newOptions: TOptions = options) => next(newOptions, index + 1),
     })
   }
 
-  return await next(options)
+  return next(options, 0)
 }