fix(server): apply error validation to event iterator

Author: unnoq

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

@@ -1,4 +1,5 @@
 import { ORPCError } from '@orpc/client'
+import { HibernationEventIterator } from '@orpc/standard-server'
 import * as z from 'zod'
 import { createORPCErrorConstructorMap, validateORPCError } from './error'
 import { isLazy, lazy, unlazy } from './lazy'
@@ -39,6 +40,16 @@ const procedure = new Procedure({
   meta: {},
 })
 
+const unvalidatedProcedure = new Procedure({
+  errorMap: baseErrors,
+  route: {},
+  handler,
+  middlewares: [preMid1, preMid2, postMid1, postMid2],
+  inputValidationIndex: 2,
+  outputValidationIndex: 2,
+  meta: {},
+})
+
 const procedureCases = [
   ['without lazy', procedure],
   ['with lazy', lazy(() => Promise.resolve({ default: procedure }))],
@@ -462,6 +473,39 @@ describe.each(procedureCases)('createProcedureClient - case %s', async (_, proce
       expect(validateORPCError).toBeCalledTimes(1)
       expect(validateORPCError).toBeCalledWith(baseErrors, e1)
     })
+
+    describe('event iterator', async () => {
+      const client = createProcedureClient(unvalidatedProcedure)
+
+      it('throw non-ORPCError right away', async () => {
+        const e1 = new Error('non-ORPC Error')
+        handler.mockImplementationOnce(async function* () {
+          throw e1
+        } as any)
+
+        const iterator = await client({ val: '123' }) as any
+
+        await expect(iterator.next()).rejects.toBe(e1)
+      })
+
+      it('validate ORPC Error', async () => {
+        const e1 = new ORPCError('BAD_REQUEST')
+        const e2 = new ORPCError('BAD_REQUEST', { defined: true })
+
+        handler.mockImplementationOnce(async function* () {
+          throw e1
+        } as any)
+        vi.mocked(validateORPCError).mockReturnValueOnce(Promise.resolve(e2))
+
+        // signal here for test coverage
+        const iterator = await client({ val: '123' }, { signal: AbortSignal.timeout(10) }) as any
+
+        await expect(iterator.next()).rejects.toBe(e2)
+
+        expect(validateORPCError).toBeCalledTimes(1)
+        expect(validateORPCError).toBeCalledWith(baseErrors, e1)
+      })
+    })
   })
 
   it('with client context', async () => {
@@ -510,6 +554,13 @@ describe.each(procedureCases)('createProcedureClient - case %s', async (_, proce
     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('not modify HibernationEventIterator', async () => {
+    const client = createProcedureClient(unvalidatedProcedure)
+    const iterator = new HibernationEventIterator(() => {})
+    handler.mockResolvedValueOnce(iterator as any)
+    await expect(client({ val: '123' })).resolves.toBe(iterator)
+  })
 })
 
 it('still work without InputSchema', async () => {

packages/server/src/procedure-client.ts

@@ -5,7 +5,7 @@ import type { Context } from './context'
 import type { ORPCErrorConstructorMap } from './error'
 import type { Lazyable } from './lazy'
 import type { AnyProcedure, Procedure, ProcedureHandlerOptions } from './procedure'
-import { ORPCError } from '@orpc/client'
+import { mapEventIterator, ORPCError } from '@orpc/client'
 import { ValidationError } from '@orpc/contract'
 import { asyncIteratorWithSpan, intercept, isAsyncIteratorObject, resolveMaybeOptionalOptions, runWithSpan, toArray, value } from '@orpc/shared'
 import { HibernationEventIterator } from '@orpc/standard-server'
@@ -98,6 +98,16 @@ export function createProcedureClient<
     const context = await value(options.context ?? {} as TInitialContext, clientContext)
     const errors = createORPCErrorConstructorMap(procedure['~orpc'].errorMap)
 
+    const validateError = async (e: unknown) => {
+      if (!(e instanceof ORPCError)) {
+        return e
+      }
+
+      const validated = await validateORPCError(procedure['~orpc'].errorMap, e)
+
+      return validated
+    }
+
     try {
       const output = await runWithSpan(
         { name: 'call_procedure', signal: callerOptions?.signal },
@@ -136,22 +146,22 @@ export function createProcedureClient<
          * If remove this return, can be breaking change
          * because AsyncIteratorClass convert `.throw` to `.return` (rarely used)
          */
-        return asyncIteratorWithSpan(
-          { name: 'consume_event_iterator_output', signal: callerOptions?.signal },
-          output,
+        return mapEventIterator(
+          asyncIteratorWithSpan(
+            { name: 'consume_event_iterator_output', signal: callerOptions?.signal },
+            output,
+          ),
+          {
+            value: v => v,
+            error: e => validateError(e),
+          },
         ) as typeof output
       }
 
       return output
     }
     catch (e) {
-      if (!(e instanceof ORPCError)) {
-        throw e
-      }
-
-      const validated = await validateORPCError(procedure['~orpc'].errorMap, e)
-
-      throw validated
+      throw await validateError(e)
     }
   }
 }