From 6f69bf04b744d88efa608dd4458dd650ad57cdf0 Mon Sep 17 00:00:00 2001
From: hayatosc <145091553+hayatosc@users.noreply.github.com>
Date: Sun, 8 Sep 2024 15:49:09 +0900
Subject: [PATCH] feat(websocket): Add generics type to `WSContext` (#3337)

* Extends WSContext's types

* Add new generics on UpgradeWebSocket types

* fix with linter

* revert createWSContext types

* prettier fix
---
 src/adapter/bun/websocket.ts                | 13 ++++++-------
 src/adapter/cloudflare-workers/websocket.ts |  4 ++--
 src/adapter/deno/websocket.ts               |  4 ++--
 src/helper/websocket/index.ts               | 20 ++++++++++----------
 4 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/src/adapter/bun/websocket.ts b/src/adapter/bun/websocket.ts
index 3f236f2b..daae7647 100644
--- a/src/adapter/bun/websocket.ts
+++ b/src/adapter/bun/websocket.ts
@@ -19,11 +19,9 @@ interface BunWebSocketHandler<T> {
   close(ws: BunServerWebSocket<T>, code?: number, reason?: string): void
   message(ws: BunServerWebSocket<T>, message: string | Uint8Array): void
 }
-interface CreateWebSocket {
-  (): {
-    upgradeWebSocket: UpgradeWebSocket
-    websocket: BunWebSocketHandler<BunWebSocketData>
-  }
+interface CreateWebSocket<T> {
+  upgradeWebSocket: UpgradeWebSocket<T>
+  websocket: BunWebSocketHandler<BunWebSocketData>
 }
 export interface BunWebSocketData {
   connId: number
@@ -49,10 +47,11 @@ const createWSContext = (ws: BunServerWebSocket<BunWebSocketData>): WSContext =>
   }
 }
 
-export const createBunWebSocket: CreateWebSocket = () => {
+export const createBunWebSocket = <T>(): CreateWebSocket<T> => {
   const websocketConns: WSEvents[] = []
 
-  const upgradeWebSocket: UpgradeWebSocket = (createEvents) => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const upgradeWebSocket: UpgradeWebSocket<any> = (createEvents) => {
     return async (c, next) => {
       const server = getBunServer(c)
       if (!server) {
diff --git a/src/adapter/cloudflare-workers/websocket.ts b/src/adapter/cloudflare-workers/websocket.ts
index c7b98265..312e3478 100644
--- a/src/adapter/cloudflare-workers/websocket.ts
+++ b/src/adapter/cloudflare-workers/websocket.ts
@@ -1,7 +1,7 @@
 import type { UpgradeWebSocket, WSContext, WSReadyState } from '../../helper/websocket'
 
 // Based on https://github.com/honojs/hono/issues/1153#issuecomment-1767321332
-export const upgradeWebSocket: UpgradeWebSocket = (createEvents) => async (c, next) => {
+export const upgradeWebSocket: UpgradeWebSocket<WebSocket> = (createEvents) => async (c, next) => {
   const events = await createEvents(c)
 
   const upgradeHeader = c.req.header('Upgrade')
@@ -14,7 +14,7 @@ export const upgradeWebSocket: UpgradeWebSocket = (createEvents) => async (c, ne
   const client: WebSocket = webSocketPair[0]
   const server: WebSocket = webSocketPair[1]
 
-  const wsContext: WSContext = {
+  const wsContext: WSContext<WebSocket> = {
     binaryType: 'arraybuffer',
     close: (code, reason) => server.close(code, reason),
     get protocol() {
diff --git a/src/adapter/deno/websocket.ts b/src/adapter/deno/websocket.ts
index 6dc3db11..1eda3e15 100644
--- a/src/adapter/deno/websocket.ts
+++ b/src/adapter/deno/websocket.ts
@@ -20,7 +20,7 @@ export interface UpgradeWebSocketOptions {
   idleTimeout?: number
 }
 
-export const upgradeWebSocket: UpgradeWebSocket<UpgradeWebSocketOptions> =
+export const upgradeWebSocket: UpgradeWebSocket<WebSocket, UpgradeWebSocketOptions> =
   (createEvents, options) => async (c, next) => {
     if (c.req.header('upgrade') !== 'websocket') {
       return await next()
@@ -29,7 +29,7 @@ export const upgradeWebSocket: UpgradeWebSocket<UpgradeWebSocketOptions> =
     const events = await createEvents(c)
     const { response, socket } = Deno.upgradeWebSocket(c.req.raw, options || {})
 
-    const wsContext: WSContext = {
+    const wsContext: WSContext<WebSocket> = {
       binaryType: 'arraybuffer',
       close: (code, reason) => socket.close(code, reason),
       get protocol() {
diff --git a/src/helper/websocket/index.ts b/src/helper/websocket/index.ts
index 512650d8..2c27b63d 100644
--- a/src/helper/websocket/index.ts
+++ b/src/helper/websocket/index.ts
@@ -10,19 +10,19 @@ import type { MiddlewareHandler } from '../../types'
 /**
  * WebSocket Event Listeners type
  */
-export interface WSEvents {
-  onOpen?: (evt: Event, ws: WSContext) => void
-  onMessage?: (evt: MessageEvent<WSMessageReceive>, ws: WSContext) => void
-  onClose?: (evt: CloseEvent, ws: WSContext) => void
-  onError?: (evt: Event, ws: WSContext) => void
+export interface WSEvents<T = unknown> {
+  onOpen?: (evt: Event, ws: WSContext<T>) => void
+  onMessage?: (evt: MessageEvent<WSMessageReceive>, ws: WSContext<T>) => void
+  onClose?: (evt: CloseEvent, ws: WSContext<T>) => void
+  onError?: (evt: Event, ws: WSContext<T>) => void
 }
 
 /**
  * Upgrade WebSocket Type
  */
-export type UpgradeWebSocket<T = any> = (
-  createEvents: (c: Context) => WSEvents | Promise<WSEvents>,
-  options?: T
+export type UpgradeWebSocket<T = unknown, U = any> = (
+  createEvents: (c: Context) => WSEvents<T> | Promise<WSEvents<T>>,
+  options?: U
 ) => MiddlewareHandler<
   any,
   string,
@@ -33,14 +33,14 @@ export type UpgradeWebSocket<T = any> = (
 
 export type WSReadyState = 0 | 1 | 2 | 3
 
-export type WSContext = {
+export type WSContext<T = unknown> = {
   send(
     source: string | ArrayBuffer | Uint8Array,
     options?: {
       compress: boolean
     }
   ): void
-  raw?: unknown
+  raw?: T
   binaryType: BinaryType
   readyState: WSReadyState
   url: URL | null