feat: add websocket log (#798)

Signed-off-by: ryjiang <jiangruiyi@gmail.com>

server/src/crons/crons.controller.ts

@@ -42,7 +42,7 @@ export class CronsController {
   }
 
   deleteCronJob(clientId: string) {
-    console.info(`Deleting all cron jobs for client: ${clientId}`);
+    // console.info(`Deleting all cron jobs for client: ${clientId}`);
     this.schedulerRegistry.deleteAllCronJobs(clientId);
   }
 }

server/src/crons/crons.service.ts

@@ -157,7 +157,7 @@ export class SchedulerRegistry {
   }
 
   deleteAllCronJobs(clientId: string) {
-    console.log('Deleting all cron jobs in service for client:', clientId);
+    // console.log('Deleting all cron jobs in service for client:', clientId);
     this.cronJobMap.forEach((v, k) => {
       if (v.clientId === clientId) {
         v.task.stop();

server/src/socket.ts

@@ -7,12 +7,47 @@ import { isElectron } from './utils';
 export let io: Server;
 export let clients = new Map<string, Socket>();
 
+export const logWebSocketRequest = (
+  socket: Socket,
+  event: string,
+  message?: any,
+  direction: 'in' | 'out' = 'in'
+) => {
+  const clientId = socket.handshake.query['milvus-client-id'] as string;
+  const remoteAddress =
+    socket.handshake.address === '::1' ? '127.0.0.1' : socket.handshake.address;
+  const timestamp = new Date().toLocaleTimeString();
+  const directionLabel = direction === 'in' ? '←' : '→';
+  const mes =
+    process.env.ATTU_LOG_LEVEL === 'debug'
+      ? JSON.stringify(message)
+      : message
+        ? JSON.stringify(message).length
+        : '';
+
+  const logMessage = [
+    chalk.blue.bold('WS'),
+    chalk.magenta.bold(event),
+    chalk.green.bold(`Client: ${clientId}`),
+    chalk.yellow(remoteAddress),
+    chalk.gray(`@ ${timestamp}`),
+    direction === 'in'
+      ? chalk.cyan(directionLabel)
+      : chalk.magenta(directionLabel),
+    message ? chalk.white(mes) : '',
+  ]
+    .filter(Boolean)
+    .join(' ');
+
+  console.log(logMessage);
+};
+
 export function initWebSocket(server: http.Server) {
   io = new Server(server, {
     cors: {
       origin: '*',
       methods: ['GET', 'POST'],
-    }
+    },
   });
 
   io.on('connection', (socket: Socket) => {
@@ -20,27 +55,37 @@ export function initWebSocket(server: http.Server) {
 
     if (clientId) {
       clients.set(clientId, socket);
-      if (!isElectron()) {
-        console.info(chalk.green(`ws client connected ${clientId}`));
-      }
+
+      // Log connection event
+      logWebSocketRequest(socket, 'CONNECT');
 
       socket.on('disconnect', () => {
-        if (!isElectron()) {
-          console.info(chalk.green(`ws client disconnected ${clientId}`));
-        }
+        // Log disconnect event
+        logWebSocketRequest(socket, 'DISCONNECT');
         clients.delete(clientId);
       });
 
       socket.on('error', (error: Error) => {
-        if (!isElectron()) {
-          console.error(
-            chalk.red(`ws client error ${clientId}: ${error.message}`)
-          );
-        }
+        // Log error event
+        logWebSocketRequest(socket, 'error', error.message);
+      });
+
+      // Log custom events
+      socket.onAny((event, ...args) => {
+        logWebSocketRequest(socket, event, args);
       });
     }
   });
 
+  io.use((socket, next) => {
+    const originalEmit = socket.emit;
+    socket.emit = function (event: string, ...args: any[]) {
+      logWebSocketRequest(socket, event, args, 'out');
+      return originalEmit.apply(socket, [event, ...args]);
+    };
+    next();
+  });
+
   // Handle server-level errors
   io.on('error', (error: Error) => {
     if (!isElectron()) {