4bnode/add-websocket.js

4bnode is a CLI-powered backend development platform with a built-in visual dashboard to generate, manage, and test Node.js/Express APIs faster.

#!/usr/bin/env node

import path from "path";
import { confirm, input } from "@inquirer/prompts";
import { getProjectRoot } from "./lib/project.js";
import { updateEnvFile } from "./lib/env.js";
import {
  addImport,
  insertBeforeRoutes,
  replaceAppListen,
} from "./lib/indexFile.js";
import {
  showHeader,
  showTaskDone,
  showError,
  showFileAction,
  installDeps,
} from "./lib/ui.js";

async function main() {
  showHeader("add-websocket", "Add native WebSocket support");

  const samePort = await confirm({
    message: "Use the server port for WebSocket?",
    default: true,
  });

  let wsPort = null;
  if (!samePort) {
    const serverPort = process.env.PORT || 3000;
    wsPort = await input({
      message: "Enter WebSocket port:",
      default: "3001",
      validate: (v) => {
        const num = parseInt(v, 10);
        if (isNaN(num) || num < 1 || num > 65535) return "Enter a valid port.";
        if (num === Number(serverPort))
          return "Cannot be the same as the server port.";
        return true;
      },
    });
  }

  console.log();

  addImport("import { WebSocketServer } from 'ws';");
  if (samePort) {
    addImport("import http from 'http';");
  }
  showFileAction("updated", "index.js");

  const websocketSetup = samePort
    ? `
const server = http.createServer(app);
const wsServer = new WebSocketServer({ server });
console.log(\`WebSocket server is running on the same port as the server (\${port})\`);`
    : `
const wsPort = process.env.WS_PORT || ${wsPort};
const wsServer = new WebSocketServer({ port: wsPort });
console.log(\`WebSocket server is running on WebSocket port \${wsPort}\`);`;

  const wsLogic = `
wsServer.on('connection', (ws) => {
  console.log('A client connected');

  ws.on('message', (message) => {
    try {
      const parsedMessage = JSON.parse(message);
      console.log('Received:', parsedMessage);

      wsServer.clients.forEach((client) => {
        if (client !== ws && client.readyState === ws.OPEN) {
          client.send(JSON.stringify(parsedMessage));
        }
      });
    } catch (error) {
      console.error('Invalid message format:', error);
    }
  });

  ws.on('close', () => {
    console.log('Client disconnected');
  });
});

app.use((req, res, next) => {
  req.wss = wsServer;
  next();
});`;

  insertBeforeRoutes(websocketSetup + "\n" + wsLogic);

  if (samePort) {
    replaceAppListen(`
server.listen(port, () => {
  const localIp = getLocalNetworkIp();
  console.log(\`App is running on http://localhost:\${port}\`);
  if (localIp) {
    console.log(
      \`App is also accessible on your local network at http://\${localIp}:\${port}\`
    );
  }
});`);
  }

  if (!samePort) {
    const projectRoot = getProjectRoot();
    updateEnvFile(path.join(projectRoot, ".env"), "WS_PORT", wsPort);
    showFileAction("updated", ".env");
  }

  await installDeps("ws");

  showTaskDone("WebSocket integration complete", [
    samePort ? "Sharing server port" : `WebSocket port: ${wsPort}`,
    "Access via req.wss in routes",
    "Auto-broadcasts to all clients",
  ]);
}

main().catch((err) => {
  if (err.name === "ExitPromptError") process.exit(0);
  showError(err.message);
  process.exit(1);
});