Improved some stuff

src/bin/index.ts

@@ -22,9 +22,6 @@ app.prepare().then(() => {
 
   const wss = new WebSocketServer({ noServer: true });
 
-  // WebSocket connections
-  wss.on("connection", (ws: WebSocket, req: IncomingMessage) => {});
-
   // Upgrade HTTP → WebSocket
   server.on("upgrade", (req, socket, head) => {
     if (!req.url) return;

src/index.ts

@@ -1,16 +1,19 @@
 import { Duplex } from "stream";
+import { WebSocketServer } from "ws";
 
 export interface NextWebsocketConfig {
   app: NextWebsocketApp;
 }
 
 export interface WsRouteHandler {
-  upgrade?: (req: Request, socket: Duplex, head: Buffer<ArrayBuffer>) => void;
+  upgrade?: (req: Request, socket: Duplex, head: NonSharedBuffer) => void;
   post?: (req: Request) => void;
+  get?: (req: Request) => void;
 }
 
 export class NextWebsocketApp {
   handlers: Record<string, WsRouteHandler> = {};
+  websocketHandlers: ((ws: WebSocketServer) => void)[] = [];
 
   on<K extends keyof WsRouteHandler>(
     method: K,
@@ -21,17 +24,22 @@ export class NextWebsocketApp {
     this.handlers[route][method] = handler;
   }
 
-  upgrade(
-    route: string,
-    handler: (req: Request, socket: Duplex, head: Buffer<ArrayBuffer>) => void,
-  ) {
+  wss(handler: (ws: WebSocketServer) => void) {
+    this.websocketHandlers.push(handler);
+  }
+
+  upgrade(route: string, handler: WsRouteHandler["upgrade"]) {
     this.on("upgrade", route, handler);
   }
 
-  post(route: string, handler: () => void) {
+  post(route: string, handler: WsRouteHandler["post"]) {
     this.on("post", route, handler);
   }
 
+  get(route: string, handler: WsRouteHandler["get"]) {
+    this.on("get", route, handler);
+  }
+
   emit<K extends keyof WsRouteHandler>(
     method: K,
     route: string,