feat: pass ws options to ws route

elysiajs · Sep 20, 2023 · 2d42995 · 2d42995
1 parent 8afed21
commit 2d42995
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -8,6 +8,7 @@ import { transformTRPCResponse, getTRPCErrorFromUnknown } from './utils'
 
 import type { TSchema } from '@sinclair/typebox'
 import type { TRPCClientIncomingRequest, TRPCOptions } from './types'
+import { ElysiaWSOptions } from 'elysia/dist/ws'
 
 export function compile<T extends TSchema>(schema: T) {
     const check = getSchemaValidator(schema, {})
@@ -37,7 +38,7 @@ const getPath = (url: string) => {
 export const trpc =
     (
         router: Router<any>,
-        { endpoint = '/trpc', ...options }: TRPCOptions = {
+        { endpoint = '/trpc', ...options }: TRPCOptions & ElysiaWSOptions<any, any, {}> = {
             endpoint: '/trpc'
         }
     ) =>
@@ -68,7 +69,9 @@ export const trpc =
         // @ts-ignore
         if (app.wsRouter)
             app.ws<any, any>(endpoint, {
+                ...options,
                 async message(ws, message) {
+                    options.message?.(ws, message)
                     const messages: TRPCClientIncomingRequest[] = Array.isArray(
                         message
                     )
@@ -187,7 +190,8 @@ export const trpc =
                         observers.set(ws.data.id.toString(), observer)
                     }
                 },
-                close(ws) {
+                close(ws, code, message) {
+                    options.close?.(ws, code, message)
                     observers.get(ws.data.id.toString())?.unsubscribe()
                     observers.delete(ws.data.id.toString())
                 }