Server Side Code:
const WebSocket = require("ws");

const wss = new WebSocket.Server({ port: 8000 });

wss.on("connection", (clientWs) => {
  console.log("Client connected");

  const url =
    "wss://gateway.ai.cloudflare.com/v1/account_id/gateway/";

 
  // Connect directly to Cloudflare WebSocket
  const cfWs = new WebSocket(url, {
    headers: {
      "cf-aig-authorization": "Bearer gateway token",
    
    },
  });

  cfWs.on("open", () => {
    console.log("Connected to Cloudflare AI Gateway");
    clientWs.send(JSON.stringify({ status: "ready" }));
  });

  // Forward client messages to Cloudflare AI Gateway
  clientWs.on("message", (message) => {
    try {
      const { prompt } = JSON.parse(message);
      console.log(prompt);
      const payload = {
        type: "universal.create",
        request: {
          eventId: `chat-${Date.now()}`,
          provider: "workers-ai",
          endpoint: "@cf/meta/llama-3.1-8b-instruct",
          headers: {
            Authorization: "Bearer wotkers-ai token",
            "Content-Type": "application/json",
          },
          query: {
            // messages: messages.map((m) => ({
            //   role: m.role,
            //   content: m.content,
            // })),
            prompt: prompt,
          },
        },
      };
      // const payload = {
      //   type: "response.create",
      //   response: { modalities: ["text"], instructions: "Tell me a joke" },
      // };
      cfWs.send(JSON.stringify(payload));
    } catch (error) {
      console.error("Error processing message:", error);
      clientWs.send(JSON.stringify({ error: "Invalid message format" }));
    }
  });

  // Forward AI responses back to the client
  cfWs.on("message", (message) => {
    if (clientWs.readyState === WebSocket.OPEN) {
      clientWs.send(message.toString());
    }
  });

  // Cleanup when either connection closes or errors
  const cleanup = () => {
    if (cfWs.readyState === WebSocket.OPEN) cfWs.close();
    if (clientWs.readyState === WebSocket.OPEN) clientWs.close();
  };

  clientWs.on("close", (event) => {
    console.log(`WebSocket closed: ${event.code} - ${event.reason}`);
    cleanup();
  });
  cfWs.on("close", (event) => {
    console.log(`WebSocket closed: ${event.code} - ${event.reason}`);
    cleanup();
  });
  clientWs.on("error", (error) => {
    console.log(error);
    cleanup();
  });
  cfWs.on("error", (error) => {
    console.log(error);
    cleanup();
  });
});

console.log("WebSocket server running on ws://localhost:8000");
Client Side Code: 
document.addEventListener("DOMContentLoaded", () => {
        const startBtn = document.getElementById("startBtn");
        const stopBtn = document.getElementById("stopBtn");
        const textInput = document.getElementById("text");

        let socket;
        setupWebSocket();
        // Start button click
        startBtn.addEventListener("click", async () => {
          const text = textInput.value.trim();
          socket.send(JSON.stringify({ prompt: text }));
        });

        // Stop button click
        stopBtn.addEventListener("click", () => {
          if (socket) {
            socket.close();
          }
        });

        function setupWebSocket() {
          // Close existing connection if any
          if (socket) {
            socket.close();
          }

          // Create new WebSocket connection

          socket = new WebSocket("ws://localhost:8000");

          socket.onopen = () => {
            console.log("WebSocket connection established");
          };

          socket.onmessage = (event) => {
            try {
              const data = JSON.parse(event.data);
              console.log(data)
              if (data.hasOwnProperty('response')) console.log(data.response?.result?.response);
            } catch (error) {
              console.log(`Error parsing message: ${error.message}`);
            }
          };

          socket.onerror = (error) => {
            console.log(`WebSocket error: ${error.message || "Unknown error"}`);
          };

          socket.onclose = (event) => {
            console.log(`WebSocket closed: ${event.code} - ${event.reason}`);
          };
        }
      });