fix(server): return correct context from middleware next (#245)

Recent changes not restrict middleware out context match with in
context, so we need reflect it in runtime.

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

## Summary by CodeRabbit

- **Tests**
- Enhanced middleware test cases with additional assertions to verify
that each middleware component produces the expected output and context.
- **Refactor**
- Improved internal context management during middleware execution for
more consistent state propagation and simplified structure.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: dinwwwh

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

@@ -178,25 +178,21 @@ describe.each(procedureCases)('createProcedureClient - case %s', async (_, proce
     })
 
     preMid2.mockImplementationOnce(({ next }) => {
-      return next({
-        context: {
-          extra2: '__extra2__',
-        },
-      })
+      return next()
     })
 
     postMid1.mockImplementationOnce(({ next }) => {
       return next({
         context: {
-          extra3: '__extra3__',
+          extra2: '__extra2__',
         },
       })
     })
 
     postMid2.mockImplementationOnce(({ next }) => {
       return next({
         context: {
-          extra4: '__extra4__',
+          extra3: '__extra3__',
         },
       })
     })
@@ -205,36 +201,38 @@ describe.each(procedureCases)('createProcedureClient - case %s', async (_, proce
 
     expect(preMid1).toBeCalledTimes(1)
     expect(preMid1).toHaveBeenCalledWith(expect.objectContaining({ context: {} }), expect.any(Object), expect.any(Function))
+    expect(await preMid1.mock.results[0]!.value).toEqual({ output: { val: 123 }, context: { extra1: '__extra1__' } })
 
     expect(preMid2).toBeCalledTimes(1)
     expect(preMid2).toHaveBeenCalledWith(expect.objectContaining({
       context: { extra1: '__extra1__' },
     }), expect.any(Object), expect.any(Function))
+    expect(await preMid2.mock.results[0]!.value).toEqual({ output: { val: 123 }, context: { } })
 
     expect(postMid1).toBeCalledTimes(1)
     expect(postMid1).toHaveBeenCalledWith(expect.objectContaining({
       context: {
         extra1: '__extra1__',
-        extra2: '__extra2__',
       },
     }), expect.any(Object), expect.any(Function))
+    expect(await postMid1.mock.results[0]!.value).toEqual({ output: { val: '123' }, context: { extra2: '__extra2__' } })
 
     expect(postMid2).toBeCalledTimes(1)
     expect(postMid2).toHaveBeenCalledWith(expect.objectContaining({
       context: {
         extra1: '__extra1__',
         extra2: '__extra2__',
-        extra3: '__extra3__',
       },
     }), expect.any(Object), expect.any(Function))
+    expect(await postMid2.mock.results[0]!.value).toEqual({ output: { val: '123' }, context: { extra3: '__extra3__' } })
 
     expect(handler).toBeCalledTimes(1)
     expect(handler).toHaveBeenCalledWith(expect.objectContaining({ context: {
       extra1: '__extra1__',
       extra2: '__extra2__',
       extra3: '__extra3__',
-      extra4: '__extra4__',
     } }))
+    expect(await handler.mock.results[0]!.value).toEqual({ val: '123' })
   })
 
   it('middleware can override context', async () => {
@@ -280,24 +278,29 @@ describe.each(procedureCases)('createProcedureClient - case %s', async (_, proce
     expect(preMid1).toHaveBeenCalledWith(expect.objectContaining({
       context: expect.objectContaining({ userId: '123' }),
     }), expect.any(Object), expect.any(Function))
+    expect(await preMid1.mock.results[0]!.value).toEqual({ output: { val: 123 }, context: { userId: '__override1__' } })
 
     expect(preMid2).toBeCalledTimes(1)
     expect(preMid2).toHaveBeenCalledWith(expect.objectContaining({
       context: expect.objectContaining({ userId: '__override1__' }),
     }), expect.any(Object), expect.any(Function))
+    expect(await preMid2.mock.results[0]!.value).toEqual({ output: { val: 123 }, context: { userId: '__override2__' } })
 
     expect(postMid1).toBeCalledTimes(1)
     expect(postMid1).toHaveBeenCalledWith(expect.objectContaining({
       context: expect.objectContaining({ userId: '__override2__' }),
     }), expect.any(Object), expect.any(Function))
+    expect(await postMid1.mock.results[0]!.value).toEqual({ output: { val: '123' }, context: { userId: '__override3__' } })
 
     expect(postMid2).toBeCalledTimes(1)
     expect(postMid2).toHaveBeenCalledWith(expect.objectContaining({
       context: expect.objectContaining({ userId: '__override3__' }),
     }), expect.any(Object), expect.any(Function))
+    expect(await postMid2.mock.results[0]!.value).toEqual({ output: { val: '123' }, context: { userId: '__override4__' } })
 
     expect(handler).toBeCalledTimes(1)
     expect(handler).toHaveBeenCalledWith(expect.objectContaining({ context: expect.objectContaining({ userId: '__override4__' }) }))
+    expect(await handler.mock.results[0]!.value).toEqual({ val: '123' })
   })
 
   const contextCases = [

packages/server/src/procedure-client.ts

@@ -1,12 +1,12 @@
 import type { Client, ClientContext } from '@orpc/client'
 import type { Interceptor, MaybeOptionalOptions, Value } from '@orpc/shared'
-import type { Context } from './context'
 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'
 import { intercept, toError, value } from '@orpc/shared'
+import { type Context, mergeCurrentContext } from './context'
 import { createORPCErrorConstructorMap, type ORPCErrorConstructorMap, validateORPCError } from './error'
 import { unlazy } from './lazy'
 import { middlewareOutputFn } from './middleware'
@@ -171,8 +171,10 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced
 
   const next: MiddlewareNextFn<any> = async (...[nextOptions]) => {
     const index = currentIndex
+    const midContext = nextOptions?.context ?? {} as any
+
     currentIndex += 1
-    currentContext = { ...currentContext, ...nextOptions?.context }
+    currentContext = mergeCurrentContext(currentContext, midContext)
 
     if (index === inputValidationIndex) {
       currentInput = await validateInput(procedure, currentInput)
@@ -181,20 +183,26 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced
     const mid = middlewares[index]
 
     const result = mid
-      ? await mid({ ...options, context: currentContext, next }, currentInput, middlewareOutputFn)
-      : { output: await procedure['~orpc'].handler({ ...options, context: currentContext, input: currentInput }), context: currentContext }
+      ? {
+          context: midContext,
+          output: (await mid({ ...options, context: currentContext, next }, currentInput, middlewareOutputFn)).output,
+        }
+      : {
+          context: midContext,
+          output: await procedure['~orpc'].handler({ ...options, context: currentContext, input: currentInput }),
+        }
 
     if (index === outputValidationIndex) {
       const validatedOutput = await validateOutput(procedure, result.output)
 
       return {
-        ...result,
+        context: result.context,
         output: validatedOutput,
       }
     }
 
     return result
   }
 
-  return (await next({})).output
+  return (await next()).output
 }