feat: updates to websocket utils implementation

Author: vveerrgg

src/__tests__/client.test.ts

@@ -15,8 +15,8 @@ import type { NostrWSMessage } from '../types';
 describe('NostrWSClient', () => {
   let client: NostrWSClient;
   let mockLogger: ReturnType<typeof createMockLogger>;
-  let messageHandler: ReturnType<typeof vi.fn>;
-  let errorHandler: ReturnType<typeof vi.fn>;
+  let messageHandler: () => Promise<void>;
+  let errorHandler: () => void;
 
   beforeEach(() => {
     vi.clearAllMocks();

src/client.ts

@@ -93,9 +93,10 @@ export class NostrWSClient {
     });
 
     ws.on('close', () => {
-      if (this.options.handlers.close) {
+      const closeHandler = this.options.handlers.close;
+      if (closeHandler) {
         process.nextTick(() => {
-          this.options.handlers.close(ws as EnhancedWebSocket);
+          closeHandler(ws as EnhancedWebSocket);
         });
       }
       this.cleanup();

src/connection-manager.ts

@@ -2,7 +2,8 @@ import { WebSocket } from 'ws';
 import { randomBytes } from 'crypto';
 import { Logger, EnhancedWebSocket, NostrWSMessage, ConnectionStats, NostrEvent, NostrMessageType } from './types';
 import { createDefaultLogger } from './logger';
-import { NostrWSError, ErrorCodes } from './error-handler.js';
+import NostrWSErrorHandler from './error-handler.js';
+import { NostrError, NostrErrorType } from './types/errors.js';
 
 interface ConnectionManagerOptions {
   logger?: Logger;
@@ -73,6 +74,13 @@ export class ConnectionManager {
     return this.connections.get(id);
   }
 
+  /**
+   * Get all active connections
+   */
+  public getConnections(): EnhancedWebSocket[] {
+    return Array.from(this.connections.values());
+  }
+
   /**
    * Ping a connection to check health
    */

src/server.ts

@@ -4,18 +4,9 @@ import type { IncomingMessage } from 'http';
 import type { NostrWSMessage, EnhancedWebSocket, NostrWSOptions } from './types/index.js';
 import { NostrErrorType } from './types/errors.js';
 import NostrWSErrorHandler from './error-handler.js';
-import ConnectionManager from './connection-manager.js';
+import { ConnectionManager } from './connection-manager.js';
 import { createDefaultLogger } from './logger.js';
 
-/**
- * Default message handlers that do nothing
- */
-const defaultHandlers: Required<NostrWSOptions['handlers']> = {
-  message: () => Promise.resolve(),
-  error: () => {},
-  close: () => {},
-};
-
 export class NostrWSServer extends EventEmitter {
   private readonly wss: WebSocketServer;
   private readonly options: Required<NostrWSOptions>;
@@ -24,13 +15,20 @@ export class NostrWSServer extends EventEmitter {
 
   constructor(options: NostrWSOptions = {}) {
     super();
+    
+    const defaultHandlers = {
+      message: () => Promise.resolve(),
+      error: () => {},
+      close: () => {},
+    } as const;
+
     this.options = {
       logger: options.logger || createDefaultLogger(),
       heartbeatInterval: options.heartbeatInterval || 30000,
       handlers: {
-        message: options.handlers?.message ?? defaultHandlers.message,
-        error: options.handlers?.error ?? defaultHandlers.error,
-        close: options.handlers?.close ?? defaultHandlers.close,
+        message: options.handlers?.message || defaultHandlers.message,
+        error: options.handlers?.error || defaultHandlers.error,
+        close: options.handlers?.close || defaultHandlers.close,
       },
       autoReconnect: options.autoReconnect ?? true,
       maxReconnectAttempts: options.maxReconnectAttempts || 5,
@@ -73,25 +71,27 @@ export class NostrWSServer extends EventEmitter {
     });
   }
 
-  private setupWebSocket(ws: EnhancedWebSocket, request: IncomingMessage): void {
+  private setupWebSocket(ws: EnhancedWebSocket & WebSocket, request: IncomingMessage): void {
     ws.on('message', async (data: Buffer) => {
       try {
         const message = JSON.parse(data.toString()) as NostrWSMessage;
         if (!Array.isArray(message) || message.length < 2) {
           throw new Error('Invalid message format');
         }
-        if (this.options.handlers.message) {
-          await this.options.handlers.message(ws, message);
+        const messageHandler = this.options.handlers.message;
+        if (messageHandler) {
+          await messageHandler(ws, message);
         }
         this.emit('message', ws, message);
       } catch (error) {
         const wsError = this.errorHandler.handleMessageError(
           error as Error,
           data.toString(),
-          ws.id
+          (ws as EnhancedWebSocket).id
         );
-        if (this.options.handlers.error) {
-          this.options.handlers.error(ws, wsError);
+        const errorHandler = this.options.handlers.error;
+        if (errorHandler) {
+          errorHandler(ws, wsError);
         }
         this.emit('error', ws, wsError);
       }
@@ -119,7 +119,7 @@ export class NostrWSServer extends EventEmitter {
 
   public broadcast(message: NostrWSMessage): void {
     const connections = this.connectionManager.getConnections();
-    connections.forEach((ws) => {
+    connections.forEach((ws: EnhancedWebSocket & WebSocket) => {
       try {
         if (ws.readyState === WebSocket.OPEN) {
           ws.send(JSON.stringify(message));
@@ -137,7 +137,7 @@ export class NostrWSServer extends EventEmitter {
   public close(): Promise<void> {
     return new Promise((resolve) => {
       const connections = this.connectionManager.getConnections();
-      connections.forEach((ws) => {
+      connections.forEach((ws: EnhancedWebSocket & WebSocket) => {
         try {
           ws.close();
         } catch (error) {