From 36172491a41dc85d2c9372a6938b25fe7a4defc7 Mon Sep 17 00:00:00 2001
From: Mauricio Siu <47042324+Siumauricio@users.noreply.github.com>
Date: Fri, 4 Apr 2025 01:55:29 -0600
Subject: [PATCH] refactor(websocket): streamline WebSocket server setup and
 client instantiation

- Removed the request validation logic from the WebSocket connection handler.
- Added a cleanup function to close the WebSocket server.
- Introduced a singleton pattern for the WebSocket client to manage connections more efficiently.
---
 apps/dokploy/server/wss/drawer-logs.ts | 16 ++++++---------
 apps/dokploy/utils/api.ts              | 27 +++++++++++++++++++-------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/apps/dokploy/server/wss/drawer-logs.ts b/apps/dokploy/server/wss/drawer-logs.ts
index 404dfeee5..dcdeaad7d 100644
--- a/apps/dokploy/server/wss/drawer-logs.ts
+++ b/apps/dokploy/server/wss/drawer-logs.ts
@@ -1,5 +1,4 @@
 import type http from "node:http";
-import { validateRequest } from "@dokploy/server/index";
 import { applyWSSHandler } from "@trpc/server/adapters/ws";
 import { WebSocketServer } from "ws";
 import { appRouter } from "../api/root";
@@ -13,11 +12,13 @@ export const setupDrawerLogsWebSocketServer = (
 		path: "/drawer-logs",
 	});
 
+	// Set up tRPC WebSocket handler
 	applyWSSHandler({
 		wss: wssTerm,
 		router: appRouter,
 		createContext: createTRPCContext as any,
 	});
+
 	server.on("upgrade", (req, socket, head) => {
 		const { pathname } = new URL(req.url || "", `http://${req.headers.host}`);
 
@@ -31,13 +32,8 @@ export const setupDrawerLogsWebSocketServer = (
 		}
 	});
 
-	wssTerm.on("connection", async (ws, req) => {
-		const _url = new URL(req.url || "", `http://${req.headers.host}`);
-		const { user, session } = await validateRequest(req);
-
-		if (!user || !session) {
-			ws.close();
-			return;
-		}
-	});
+	// Return cleanup function
+	return () => {
+		wssTerm.close();
+	};
 };
diff --git a/apps/dokploy/utils/api.ts b/apps/dokploy/utils/api.ts
index 56197528b..7c003f481 100644
--- a/apps/dokploy/utils/api.ts
+++ b/apps/dokploy/utils/api.ts
@@ -27,15 +27,28 @@ const getWsUrl = () => {
 	const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
 	const host = window.location.host;
 
+	// Use the base URL for all tRPC WebSocket connections
 	return `${protocol}${host}/drawer-logs`;
 };
 
-const wsClient =
-	typeof window !== "undefined"
-		? createWSClient({
-				url: getWsUrl() || "",
-			})
-		: null;
+// Singleton WebSocket client instance
+let wsClientInstance: ReturnType<typeof createWSClient> | null = null;
+
+const getWsClient = () => {
+	if (typeof window === "undefined") return null;
+
+	if (!wsClientInstance) {
+		wsClientInstance = createWSClient({
+			url: getWsUrl() || "",
+			onClose: () => {
+				// Reset the instance when connection closes so it can be recreated
+				wsClientInstance = null;
+			},
+		});
+	}
+
+	return wsClientInstance;
+};
 
 /** A set of type-safe react-query hooks for your tRPC API. */
 export const api = createTRPCNext<AppRouter>({
@@ -57,7 +70,7 @@ export const api = createTRPCNext<AppRouter>({
 				splitLink({
 					condition: (op) => op.type === "subscription",
 					true: wsLink({
-						client: wsClient!,
+						client: getWsClient()!,
 					}),
 					false: splitLink({
 						condition: (op) => op.input instanceof FormData,