feat: onAsyncIteratorObjectError, onReadableStreamError

packages/client/src/async-iterator-object.test.ts

@@ -69,7 +69,7 @@ describe('wrapAsyncIteratorPreservingEventMeta', () => {
   it('preserves metadata when mapping errors', async () => {
     const error = withEventMeta(new Error('TEST'), { id: 'error-1' })
     const onError = vi.fn()
-    const mapError = vi.fn(async cause => ({ mapped: cause }))
+    const mapError = vi.fn(async cause => ({ mapped: cause } as any))
 
     const iterator = (async function* () {
       throw error
@@ -107,9 +107,9 @@ describe('wrapAsyncIteratorPreservingEventMeta', () => {
     const primitiveError = wrapAsyncIteratorPreservingEventMeta((async function* () {
       throw error
     })(), {
-      mapError: async () => 'mapped-error',
+      mapError: async () => 'mapped error' as any,
     })
 
-    await expect(primitiveError.next()).rejects.toBe('mapped-error')
+    await expect(primitiveError.next()).rejects.toBe('mapped error')
   })
 })

packages/client/src/index.ts

@@ -42,8 +42,10 @@ export {
    * @deprecated Use `asyncIteratorToUnproxiedDataStream` instead.
    */
   asyncIteratorToUnproxiedDataStream as eventIteratorToUnproxiedDataStream,
+  onAsyncIteratorObjectError,
   onError,
   onFinish,
+  onReadableStreamError,
   onStart,
   onSuccess,
   streamToAsyncIteratorObject,

packages/server/src/index.ts

@@ -98,8 +98,10 @@ export {
    * @deprecated Use `asyncIteratorToUnproxiedDataStream` instead.
    */
   asyncIteratorToUnproxiedDataStream as eventIteratorToUnproxiedDataStream,
+  onAsyncIteratorObjectError,
   onError,
   onFinish,
+  onReadableStreamError,
   onStart,
   onSuccess,
   streamToAsyncIteratorObject,

packages/server/src/procedure-client.ts

@@ -80,7 +80,7 @@ export function createProcedureClient<
     const context = await value(options.context, clientContext) as TInitialContext | undefined ?? {} as TInitialContext
     const errors = createORPCErrorConstructorMap(procedure['~orpc'].errorMap)
 
-    const reconcileError = async (e: unknown) => {
+    const reconcileError = async (e: ThrowableError) => {
       if (e instanceof ORPCError) {
         return await reconcileORPCError(procedure['~orpc'].errorMap, e)
       }
@@ -142,7 +142,7 @@ export function createProcedureClient<
        * Defined errors take priority over inferable errors.
        * `reconcileError` attempts to mark the error as defined, or keeps it inferable if that's not possible.
        */
-      throw await reconcileError(e)
+      throw await reconcileError(e as ThrowableError)
     }
   }
 }

packages/shared/src/interceptor.test-d.ts

@@ -1,8 +1,8 @@
 import type { Interceptor } from './interceptor'
 import type { PromiseWithError } from './types'
-import { onError, onFinish, onStart, onSuccess } from './interceptor'
+import { onAsyncIteratorObjectError, onError, onFinish, onReadableStreamError, onStart, onSuccess } from './interceptor'
 
-it('onStart, onSuccess, onError, onFinish can be used as an interceptor', () => {
+it('onStart, onSuccess, onError, onFinish, onAsyncIteratorObjectError, onReadableStreamError can be used as an interceptor', () => {
   const interceptors: Interceptor<{ foo: string }, PromiseWithError<'success', 'error'>>[] = []
 
   interceptors.push(onStart((options) => {
@@ -41,4 +41,20 @@ it('onStart, onSuccess, onError, onFinish can be used as an interceptor', () =>
     expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
     expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
   }))
+
+  interceptors.push(onAsyncIteratorObjectError((error, options) => {
+    expectTypeOf(error).toEqualTypeOf<'error' | Error>()
+
+    expectTypeOf(options.foo).toEqualTypeOf<string>()
+    expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
+    expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
+  }))
+
+  interceptors.push(onReadableStreamError((error, options) => {
+    expectTypeOf(error).toEqualTypeOf<'error' | Error>()
+
+    expectTypeOf(options.foo).toEqualTypeOf<string>()
+    expectTypeOf(options.next).toBeCallableWith<[options?: { foo: string }]>()
+    expectTypeOf(options.next()).toEqualTypeOf<PromiseWithError<'success', 'error'>>()
+  }))
 })

packages/shared/src/interceptor.test.ts

@@ -1,4 +1,4 @@
-import { intercept, onError, onFinish, onStart, onSuccess } from './interceptor'
+import { intercept, onAsyncIteratorObjectError, onError, onFinish, onReadableStreamError, onStart, onSuccess } from './interceptor'
 
 describe('intercept', () => {
   const interceptor1 = vi.fn(({ next }) => next())
@@ -381,3 +381,133 @@ describe('lifecycle interceptors', () => {
     expect(onSuccessFn).toHaveBeenCalledTimes(0)
   })
 })
+
+describe('stream interceptors', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  describe('onAsyncIteratorObjectError', () => {
+    it('does not invoke callback for non-iterator output', async () => {
+      const callback = vi.fn()
+
+      const result = await intercept(
+        [onAsyncIteratorObjectError(callback)],
+        { foo: 'bar' },
+        () => '__OUTPUT__',
+      )
+
+      expect(result).toBe('__OUTPUT__')
+      expect(callback).not.toHaveBeenCalled()
+    })
+
+    it('invokes callback when the returned async iterator errors while being consumed', async () => {
+      const error = new Error('__ITERATOR_ERROR__')
+      const iterator = (async function* () {
+        yield 'event'
+        throw error
+      })()
+
+      const callback = vi.fn()
+
+      const wrapped = await intercept(
+        [onAsyncIteratorObjectError(callback)],
+        { foo: 'bar' },
+        () => iterator,
+      ) as AsyncIterableIterator<string>
+
+      await expect(wrapped.next()).resolves.toEqual({ value: 'event', done: false })
+      await expect(wrapped.next()).rejects.toThrow(error)
+
+      expect(callback).toHaveBeenCalledTimes(1)
+      expect(callback).toHaveBeenCalledWith(
+        error,
+        expect.objectContaining({ foo: 'bar', next: expect.any(Function) }),
+      )
+    })
+
+    it('does not invoke callback when the async iterator completes successfully', async () => {
+      const iterator = (async function* () {
+        yield 'event'
+        return 'done'
+      })()
+
+      const callback = vi.fn()
+
+      const wrapped = await intercept(
+        [onAsyncIteratorObjectError(callback)],
+        { foo: 'bar' },
+        () => iterator,
+      ) as AsyncIterableIterator<string>
+
+      await expect(wrapped.next()).resolves.toEqual({ value: 'event', done: false })
+      await expect(wrapped.next()).resolves.toEqual({ value: 'done', done: true })
+
+      expect(callback).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('onReadableStreamError', () => {
+    it('does not invoke callback for non-stream output', async () => {
+      const callback = vi.fn()
+
+      const result = await intercept(
+        [onReadableStreamError(callback)],
+        { foo: 'bar' },
+        () => '__OUTPUT__',
+      )
+
+      expect(result).toBe('__OUTPUT__')
+      expect(callback).not.toHaveBeenCalled()
+    })
+
+    it('invokes callback when the returned readable stream errors while being consumed', async () => {
+      const error = new Error('__STREAM_ERROR__')
+      const stream = new ReadableStream({
+        pull(controller) {
+          controller.error(error)
+        },
+      })
+
+      const callback = vi.fn()
+
+      const wrapped = await intercept(
+        [onReadableStreamError(callback)],
+        { foo: 'bar' },
+        () => stream,
+      ) as ReadableStream<string>
+
+      const reader = wrapped.getReader()
+      await expect(reader.read()).rejects.toThrow(error)
+
+      expect(callback).toHaveBeenCalledTimes(1)
+      expect(callback).toHaveBeenCalledWith(
+        error,
+        expect.objectContaining({ foo: 'bar', next: expect.any(Function) }),
+      )
+    })
+
+    it('does not invoke callback when the readable stream completes successfully', async () => {
+      const stream = new ReadableStream({
+        start(controller) {
+          controller.enqueue('event')
+          controller.close()
+        },
+      })
+
+      const callback = vi.fn()
+
+      const wrapped = await intercept(
+        [onReadableStreamError(callback)],
+        { foo: 'bar' },
+        () => stream,
+      ) as ReadableStream<string>
+
+      const reader = wrapped.getReader()
+      await expect(reader.read()).resolves.toEqual({ value: 'event', done: false })
+      await expect(reader.read()).resolves.toEqual({ value: undefined, done: true })
+
+      expect(callback).not.toHaveBeenCalled()
+    })
+  })
+})

packages/shared/src/interceptor.ts

@@ -1,5 +1,9 @@
 import type { Promisable } from 'type-fest'
 import type { PromiseWithError, ThrowableError } from './types'
+import { isAsyncIteratorObject } from '@standardserver/shared'
+import { wrapAsyncIterator } from './iterator'
+import { override } from './proxy'
+import { wrapReadableStream } from './stream'
 
 export type InterceptableOptions = Record<string, any>
 
@@ -96,6 +100,92 @@ export function onFinish<T, TOptions extends { next: () => any }, TRest extends
   }
 }
 
+/**
+ * Creates a middleware or interceptor that invokes a callback when the returned async
+ * iterator object throws an error while being consumed.
+ *
+ * This does not replace the `onError`. `onError` only fires on the
+ * initial call (before the interceptor returns the iterator), whereas this
+ * callback only fires while consuming the iterator. Use both together to
+ * catch all possible errors.
+ */
+export function onAsyncIteratorObjectError<
+  T,
+  TOptions extends { next: () => any },
+  TRest extends any[],
+>(
+  callback: NoInfer<(
+    error: ThrowableError | (ReturnType<TOptions['next']> extends PromiseWithError<any, infer E> ? E : ThrowableError),
+    options: TOptions,
+    ...rest: TRest
+  ) => Promisable<void>>,
+): (options: TOptions, ...rest: TRest) => T | Promise<Awaited<ReturnType<TOptions['next']>>> {
+  // The typed error should always be combined with `ThrowableError`
+  // because an AsyncIterator can throw any error during iteration,
+  // while TError only represents errors from the initial promise.
+  // In oRPC client/server usage, initial and iteration errors are usually the same,
+  // but this utility is shared, so it needs to support the general case.
+
+  return async (options, ...rest) => {
+    const output = await options.next()
+
+    if (!isAsyncIteratorObject(output)) {
+      return output
+    }
+
+    /**
+     * @warning
+     * Remember use `override` for AsyncIteratorObject to remain other special properties
+     */
+    return override(output, wrapAsyncIterator(output, {
+      onError: error => callback(error, options, ...rest),
+    }))
+  }
+}
+
+/**
+ * Creates an interceptor that invokes a callback when the returned readable
+ * stream errors while being consumed.
+ *
+ * This does not replace the `onError`. `onError` only fires on the
+ * initial call (before the interceptor returns the stream), whereas this
+ * callback only fires while consuming the stream. Use both together to catch
+ * all possible errors.
+ */
+export function onReadableStreamError<
+  T,
+  TOptions extends { next: () => any },
+  TRest extends any[],
+>(
+  callback: NoInfer<(
+    error: ThrowableError | (ReturnType<TOptions['next']> extends PromiseWithError<any, infer E> ? E : ThrowableError),
+    options: TOptions,
+    ...rest: TRest
+  ) => Promisable<void>>,
+): (options: TOptions, ...rest: TRest) => T | Promise<Awaited<ReturnType<TOptions['next']>>> {
+  // The typed error should always be combined with `ThrowableError`
+  // because a ReadableStream can throw any error during iteration,
+  // while TError only represents errors from the initial promise.
+  // In oRPC client/server usage, initial and iteration errors are usually the same,
+  // but this utility is shared, so it needs to support the general case.
+
+  return async (options, ...rest) => {
+    const output = await options.next()
+
+    if (!(output instanceof ReadableStream)) {
+      return output
+    }
+
+    /**
+     * @warning
+     * Remember use `override` for ReadableStream to remain other special properties
+     */
+    return override(output, wrapReadableStream(output, {
+      onError: error => callback(error, options, ...rest),
+    }))
+  }
+}
+
 export function intercept<TOptions extends InterceptableOptions, TResult>(
   interceptors: undefined | Interceptor<TOptions, TResult>[],
   options: NoInfer<TOptions>,

packages/shared/src/iterator.ts

@@ -14,8 +14,8 @@ export interface WrapAsyncIteratorOptions<TYield, TReturn, TMappedYield, TMapped
    */
   runWith?: <T>(run: () => Promise<T>) => Promise<T>
   mapResult?: (result: IteratorResult<TYield, TReturn>) => Promisable<IteratorResult<TMappedYield, TMappedReturn>>
-  mapError?: (error: unknown) => Promisable<unknown>
-  onError?: (error: unknown) => Promisable<void>
+  mapError?: (error: ThrowableError) => Promisable<ThrowableError>
+  onError?: (error: ThrowableError) => Promisable<void>
 
   /**
    * Execute after the stream finishes or is cancelled.
@@ -46,8 +46,8 @@ export function wrapAsyncIterator<TYield, TReturn, TMappedYield = TYield, TMappe
       return mapResult ? await mapResult(result) : result as any
     }
     catch (error) {
-      await onError?.(error)
-      throw mapError ? await mapError(error) : error
+      await onError?.(error as ThrowableError)
+      throw mapError ? await mapError(error as ThrowableError) : error
     }
   }, async (_state) => {
     try {
@@ -58,7 +58,7 @@ export function wrapAsyncIterator<TYield, TReturn, TMappedYield = TYield, TMappe
           await runWith(async () => iterator.return?.())
         }
         catch (error) {
-          await onError?.(error)
+          await onError?.(error as ThrowableError)
           throw error
         }
       }