Merge pull request #45 from AxiomChat/node-dm-patch

Node dm patch

src/app/auth/CreateAccount.tsx

@@ -83,6 +83,7 @@ export default function CreateAccount() {
         Cookies.set("verify_temp_password", password);
         router.push(`/verify`);
         Cookies.set("servers", "");
+        localStorage.clear();
       } else {
         setFeedback({ message: response.message, kind: "error" });
       }

src/app/auth/Login.tsx

@@ -30,6 +30,7 @@ export default function Login() {
       });
 
       Cookies.set("servers", "");
+      localStorage.clear();
       router.push("/chat");
     } catch (err: any) {
       setFeedback({

src/app/chat/StartNewDM.tsx

@@ -13,16 +13,20 @@ import { PlusCircle } from "lucide-react";
 import { UserProfile } from "@/hooks/get-user";
 import { useRouter } from "next/navigation";
 import axios from "axios";
+import { getProfileByUsername } from "@/actions/get-profile";
+import { toast } from "sonner";
 
 export default function StartNewDM() {
   const [username, setUsername] = useState("");
   const router = useRouter();
 
   const handleStartDM = async () => {
-    const res = await axios.get("/api/profile", {
+    const user = await getProfileByUsername(username);
-      params: { username },
+    if (!user) {
-    });
+      toast.error("User not found");
-    const user = res.data as UserProfile;
+      return;
+    }
+
     router.push(`/chat/${user.id}`);
     setUsername("");
   };

src/app/chat/[id]/page.tsx

@@ -4,45 +4,72 @@ import MessageBox from "@/components/channel/MessageBox";
 import AppLayout from "@/components/app/AppLayout";
 import { useParams, useRouter } from "next/navigation";
 import useApp from "@/hooks/use-app";
-import { useEffect, useRef, useState } from "react";
+import { useEffect } from "react";
-import auth from "@/lib/auth";
 import useAsync from "@/hooks/use-async";
 import { Button } from "@/components/ui/button";
 import useIndicators from "@/hooks/use-indicators";
+import { ClientMessage } from "@/types/protocol";
+import { getNode, useNullNode } from "@/hooks/use-node";
 import { useEffectOnceWhenReady } from "@/hooks/use-once";
 
 export default function DMs() {
   const { id } = useParams<{ id: string }>();
   const app = useApp();
-  const targetNode = useRef<WebSocket | null>(null);
   const { value: target, loading } = useAsync(() => app.getUserById(id));
   const router = useRouter();
   const { indicators, addIndicator } = useIndicators();
 
+  const sendRequest = async (req: ClientMessage) => {
+    console.log("Sending request:", req);
+
+    await targetNode.current?.waitUntilReady();
+    await app.node.current?.waitUntilReady();
+
+    switch (req.type) {
+      case "send_message":
+        targetNode.current?.socket?.send(JSON.stringify(req));
+        if (target?.node_address !== app.profile?.node_address)
+          app.node.current?.socket?.send(JSON.stringify(req));
+        break;
+      case "load_chunk":
+        targetNode.current?.socket?.send(JSON.stringify(req));
+        if (target?.node_address !== app.profile?.node_address)
+          app.node.current?.socket?.send(JSON.stringify(req));
+        break;
+      default:
+        app.node.current?.socket?.send(JSON.stringify(req));
+        break;
+    }
+  };
+
+  const targetNode = useNullNode();
+
   useEffectOnceWhenReady(
     () => {
-      async function connectTargetNode() {
+      if (!target) return;
-        if (!target) return;
+
-        auth({
+      return getNode(
-          id: target.node_address,
+        target?.node_address,
-          wsRef: targetNode,
+        app.privateMessages,
-          messageStore: app.privateMessages,
+        {
           onIndicator: addIndicator,
-        });
+          setServer: () => targetNode.current?.loadInitialMessages(id),
-      }
+        },
-
+        targetNode
-      connectTargetNode();
+      );
-
-      app.setCurrentChannel({
-        id: id,
-        name: target?.display_name || id,
-        kind: "text",
-      });
     },
     [target],
     [(v) => v]
   );
 
+  useEffect(() => {
+    app.setCurrentChannel({
+      id: id,
+      name: target?.display_name || id,
+      kind: "text",
+    });
+  }, [targetNode]);
+
   return loading || target ? (
     <AppLayout app={app}>
       <MessageBox
@@ -54,25 +81,7 @@ export default function DMs() {
             (m.channel_id === id && m.from === app.profile?.id) ||
             (m.from === id && m.channel_id === app.profile?.id)
         )}
-        sendRequest={(req) => {
+        sendRequest={sendRequest}
-          console.log("Sending request: ", req);
-          while (
-            targetNode.current?.CONNECTING ||
-            app.node.current?.CONNECTING
-          ) {}
-          switch (req.type) {
-            case "send_message":
-              targetNode.current?.send(JSON.stringify(req));
-              break;
-            case "load_chunk":
-              targetNode.current?.send(JSON.stringify(req));
-              app.node.current?.send(JSON.stringify(req));
-              break;
-            default:
-              app.node.current?.send(JSON.stringify(req));
-              break;
-          }
-        }}
         indicators={indicators.filter(
           (indicator) => indicator.indicator.params.user_id !== app.profile?.id
         )}

src/app/chat/page.tsx

@@ -33,7 +33,6 @@ export default function ChatHub() {
 
   const app = useApp();
   const recentPeople = app.dms.slice(0, 3);
-  const activeServers: Server[] = Object.values(app.servers).slice(0, 3);
   const trendingServers: Server[] = [
     {
       id: "7de3138f-e0b6-4d96-8158-8c7e866ba0ff",

src/app/server/[id]/page.tsx

@@ -0,0 +1,68 @@
+"use client";
+
+import { useEffect } from "react";
+import { useParams, useSearchParams } from "next/navigation";
+import AppLayout from "@/components/app/AppLayout";
+import useMessages from "@/hooks/use-messages";
+import useApp from "@/hooks/use-app";
+import ChannelView from "@/components/channel/ChannelView";
+import useIndicators from "@/hooks/use-indicators";
+import useNode from "@/hooks/use-node";
+
+export default function Server() {
+  const { id } = useParams<{ id: string }>();
+  const searchParams = useSearchParams();
+  const app = useApp();
+  const messageStore = useMessages();
+  const { indicators, addIndicator } = useIndicators();
+
+  const serverNode = useNode(id, messageStore, {
+    setServer: app.setServer,
+    onIndicator: addIndicator,
+    onVoiceJoin: (userId, channelId, voiceId) =>
+      app.setVoiceConns((prev) => ({
+        ...prev,
+        [channelId]: { ...prev[channelId], [userId]: voiceId },
+      })),
+    onVoiceLeave: (userId, channelId) =>
+      app.setVoiceConns((prev) => {
+        const { [userId]: _, ...remainingUsers } = prev[channelId] || {};
+        return {
+          ...prev,
+          [channelId]: remainingUsers,
+        };
+      }),
+    onSpeaking: (voiceId) =>
+      app.setSpeaking((speaking) => ({
+        ...speaking,
+        [voiceId]: Date.now(),
+      })),
+  });
+
+  useEffect(() => {
+    const channelId = searchParams.get("ch");
+    if (channelId) {
+      app.setCurrentChannel(
+        app.server?.channels.find((v) => v.id === channelId) ||
+          app.server?.channels[0] ||
+          null
+      );
+    } else {
+      app.setCurrentChannel(app.server?.channels[0] || null);
+    }
+  }, [app.server]);
+
+  return (
+    <AppLayout app={app}>
+      <ChannelView
+        app={app}
+        node={serverNode}
+        messages={messageStore.messages.filter(
+          (m) => m.channel_id === app.currentChannel?.id
+        )}
+        ip={id}
+        indicators={indicators}
+      />
+    </AppLayout>
+  );
+}

src/app/server/[ip]/page.tsx

@@ -1,81 +0,0 @@
-"use client";
-
-import { useEffect, useRef, useState } from "react";
-import { useParams, useSearchParams } from "next/navigation";
-import AppLayout from "@/components/app/AppLayout";
-import auth from "@/lib/auth";
-import useMessages from "@/hooks/use-messages";
-import useApp from "@/hooks/use-app";
-import { useEffectOnceWhenReady } from "@/hooks/use-once";
-import ChannelView from "@/components/channel/ChannelView";
-import { IndicatorContext } from "@/types/protocol";
-import useIndicators from "@/hooks/use-indicators";
-
-export default function Server() {
-  const { ip } = useParams<{ ip: string }>();
-  const searchParams = useSearchParams();
-  const serverRef = useRef<WebSocket | null>(null);
-  const app = useApp();
-  const messageStore = useMessages();
-  const { indicators, addIndicator } = useIndicators();
-
-  useEffectOnceWhenReady(
-    () => {
-      if (!ip) return;
-      auth({
-        id: ip,
-        wsRef: serverRef,
-        setServer: app.setServer,
-        messageStore: messageStore,
-        onIndicator: addIndicator,
-        onVoiceJoin: (userId, channelId, voiceId) =>
-          app.setVoiceConns((prev) => ({
-            ...prev,
-            [channelId]: { ...prev[channelId], [userId]: voiceId },
-          })),
-        onVoiceLeave: (userId, channelId) =>
-          app.setVoiceConns((prev) => {
-            const { [userId]: _, ...remainingUsers } = prev[channelId] || {};
-            return {
-              ...prev,
-              [channelId]: remainingUsers,
-            };
-          }),
-        onSpeaking: (voiceId) =>
-          app.setSpeaking((speaking) => ({
-            ...speaking,
-            [voiceId]: Date.now(),
-          })),
-      });
-    },
-    [ip, messageStore.messages],
-    [undefined, (v) => v]
-  );
-
-  useEffect(() => {
-    const channelId = searchParams.get("ch");
-    if (channelId) {
-      app.setCurrentChannel(
-        app.server?.channels.find((v) => v.id === channelId) ||
-          app.server?.channels[0] ||
-          null
-      );
-    } else {
-      app.setCurrentChannel(app.server?.channels[0] || null);
-    }
-  }, [app.server]);
-
-  return (
-    <AppLayout app={app}>
-      <ChannelView
-        app={app}
-        messages={messageStore.messages.filter(
-          (m) => m.channel_id === app.currentChannel?.id
-        )}
-        ip={ip}
-        serverRef={serverRef}
-        indicators={indicators}
-      />
-    </AppLayout>
-  );
-}

src/components/app/AppLayout.tsx

@@ -4,7 +4,6 @@ import React, { useEffect } from "react";
 import { cn } from "@/lib/utils";
 import { useRouter } from "next/navigation";
 import { toast } from "sonner";
-import auth from "@/lib/auth";
 import { useIsMobile } from "@/hooks/is-mobile";
 import App from "@/types/app";
 import { Message } from "@/types/types";
@@ -14,6 +13,8 @@ import SidebarDMs from "./SidebarDMs";
 import SidebarChannels from "./SidebarChannels";
 import DMItem from "./DMItem";
 import { useEffectOnceWhenReady } from "@/hooks/use-once";
+import { getNode, useNullNode } from "@/hooks/use-node";
+import { ClientMessage } from "@/types/protocol";
 
 export default function AppLayout({
   children,
@@ -28,44 +29,48 @@ export default function AppLayout({
 }) {
   const router = useRouter();
   const isMobile = useIsMobile();
+  const node = useNullNode();
 
   useEffectOnceWhenReady(
     () => {
       if (!app.profile?.node_address) return;
 
-      auth({
+      return getNode(
-        id: app.profile.node_address,
+        app.profile?.node_address,
-        wsRef: app.node,
+        app.privateMessages,
-        onNewMessage: (m) => {
+        {
-          if (m.from !== chatWith && m.from !== app.profile?.id)
+          onNewMessage: (m) => {
-            toast(
+            if (m.from !== chatWith && m.from !== app.profile?.id)
-              `New message from ${
+              toast(
-                app.profiles[m.from]?.display_name || m.from
+                `New message from ${
-              }`,
+                  app.profiles[m.from]?.display_name || m.from
-              {
+                }`,
-                description: m.contents.slice(0, 80),
+                {
-                action: {
+                  description: m.contents.slice(0, 80),
-                  label: "View",
+                  action: {
-                  onClick: () => router.push(`/chat/${m.from}`),
+                    label: "View",
-                },
+                    onClick: () => router.push(`/chat/${m.from}`),
-                position: "top-right",
+                  },
-              }
+                  position: "top-right",
-            );
+                }
+              );
 
-          if (onNewMessage) onNewMessage(m);
+            if (onNewMessage) onNewMessage(m);
+          },
+          setServer: () =>
+            chatWith && app.node.current?.loadInitialMessages(chatWith),
         },
-        messageStore: app.privateMessages,
+        node
-      });
+      );
-
-      app.setProfiles((prev) => {
-        if (app.profile) prev[app.profile.id] = app.profile;
-        return prev;
-      });
     },
     [chatWith, app.profile],
     [undefined, undefined]
   );
 
+  useEffect(() => {
+    app.node.current = node.current;
+  }, [node]);
+
   useEffect(() => app.setSidebarOpen(!isMobile), [isMobile]);
 
   return (

src/components/app/SidebarDMs.tsx

@@ -10,6 +10,8 @@ import DMItem from "./DMItem";
 import App from "@/types/app";
 import { UserProfile } from "@/hooks/get-user";
 import { useEffect, useState } from "react";
+import { toast } from "sonner";
+import { getProfileByUsername } from "@/actions/get-profile";
 
 export default function SidebarDMs({ app }: { app: App }) {
   const router = useRouter();
@@ -78,11 +80,14 @@ export default function SidebarDMs({ app }: { app: App }) {
                 variant="ghost"
                 className="w-full justify-start gap-2 text-sm truncate"
                 onClick={async () => {
-                  const res = await axios.get("/api/profile", {
+                  const user = await getProfileByUsername(query);
-                    params: { username: query },
+                  if (!user) {
-                  });
+                    toast.error("User not found");
-                  const user = res.data as UserProfile;
+                    return;
+                  }
+
                   router.push(`/chat/${user.id}`);
+                  setQuery("");
                 }}
               >
                 <Plus className="h-4 w-4" />

src/components/channel/ChannelView.tsx

@@ -1,29 +1,37 @@
 import { Message } from "@/types/types";
 import MessageBox from "./MessageBox";
 import App from "@/types/app";
-import React from "react";
+import React, { RefObject } from "react";
 import PluginChannel from "./PluginChannel";
-import { IndicatorContext } from "@/types/protocol";
+import { ClientMessage, IndicatorContext } from "@/types/protocol";
 import VoiceBox from "./voice/VoiceBox";
 import { sendVoice } from "@/lib/audio";
+import { NodeClient } from "@/lib/node";
 
 export default function ChannelView({
   app,
   messages,
   ip,
-  serverRef,
+  node,
   indicators,
 }: {
   app: App;
   messages: Message[];
   ip: string;
-  serverRef: React.RefObject<WebSocket | null>;
   indicators: IndicatorContext[];
+  node: RefObject<NodeClient | null>;
 }) {
   if (!app.currentChannel)
     return (
       <div className="h-svh w-full max-h-svh flex flex-col pb-5 pl-5 gap-5" />
     );
+
+  const sendRequest = async (req: ClientMessage) => {
+    await node.current?.waitUntilReady();
+    node.current?.socket?.send(JSON.stringify(req));
+    console.log(node.current?.socket);
+  };
+
   switch (app.currentChannel.kind) {
     case "text":
       return (
@@ -31,10 +39,7 @@ export default function ChannelView({
           app={app}
           channelName={app.currentChannel.name}
           messages={messages}
-          sendRequest={(req) => {
+          sendRequest={sendRequest}
-            while (serverRef.current?.CONNECTING) {}
-            serverRef.current?.send(JSON.stringify(req));
-          }}
           indicators={indicators.filter(
             (indicator) =>
               indicator.indicator.params.user_id !== app.profile?.id
@@ -49,13 +54,13 @@ export default function ChannelView({
           channelName={app.currentChannel.name}
           channelId={app.currentChannel.id}
           messages={messages}
-          sendRequest={(req) => {
+          sendRequest={sendRequest}
-            while (serverRef.current?.CONNECTING) {}
-            serverRef.current?.send(JSON.stringify(req));
-          }}
           sendVoice={(data) => {
-            while (!serverRef.current || serverRef.current?.CONNECTING) {}
+            if (!node.current?.socket) {
-            sendVoice(serverRef.current, data);
+              return;
+            }
+
+            sendVoice(node.current?.socket, data);
           }}
           indicators={indicators.filter(
             (indicator) =>
@@ -66,12 +71,7 @@ export default function ChannelView({
 
     default:
       return (
-        <PluginChannel
+        <PluginChannel app={app} messages={messages} ip={ip} node={node} />
-          app={app}
-          messages={messages}
-          ip={ip}
-          serverRef={serverRef}
-        />
       );
   }
 }

src/components/channel/PluginChannel.tsx

@@ -2,17 +2,18 @@ import { Message } from "@/types/types";
 import App from "@/types/app";
 import React, { useEffect, useRef } from "react";
 import { useParentCommunication } from "react-iframes-bridge";
+import { NodeClient } from "@/lib/node";
 
 export default function PluginChannel({
   app,
   messages,
   ip,
-  serverRef,
+  node,
 }: {
   app: App;
   messages: Message[];
   ip: string;
-  serverRef: React.RefObject<WebSocket | null>;
+  node: React.RefObject<NodeClient | null>;
 }) {
   if (!app.currentChannel)
     return (
@@ -31,8 +32,9 @@ export default function PluginChannel({
       parentComm.sendToChild("MESSAGES", messages);
     });
 
-    parentComm.onMessage("SEND_MESSAGE", (msg: string) => {
+    parentComm.onMessage("SEND_MESSAGE", async (msg: string) => {
-      serverRef.current?.send(
+      await node.current?.waitUntilReady();
+      node.current?.socket?.send(
         JSON.stringify({
           type: "send_message",
           params: {

src/hooks/use-app.ts

@@ -9,9 +9,10 @@ import Cookies from "js-cookie";
 import { toast } from "sonner";
 import { getProfileById } from "@/actions/get-profile";
 import { getServerById } from "@/actions/get-server";
+import { NodeClient } from "@/lib/node";
 
 export default function useApp(): App {
-  const node = useRef<WebSocket | null>(null);
+  const node = useRef<NodeClient | null>(null);
   const [profiles, setProfiles] = useState<StringMap<UserProfile | null>>({});
   const [sidebarOpen, setSidebarOpen] = useState(false);
   const profile = useUser();

src/hooks/use-messages.ts

@@ -10,23 +10,45 @@ export interface MessageStore {
   clearMessages: () => void;
   insertMessages: (msgs: Message[]) => void;
 }
+
 export default function useMessages(): MessageStore {
-  const [messages, setMessages] = useState<Message[]>([]);
+  const [messages, setMessagesState] = useState<Message[]>([]);
+
+  // Helper to ensure unique messages by id
+  const uniqueMessages = (msgs: Message[]) => {
+    const map = new Map<number, Message>();
+    msgs.forEach((m) => map.set(m.id, m)); // duplicates overwrite
+    return Array.from(map.values());
+  };
+
+  const setMessages = (msgs: Message[]) =>
+    setMessagesState(uniqueMessages(msgs));
+
+  const addMessage = (msg: Message) =>
+    setMessagesState((messages) =>
+      messages.some((m) => m.id === msg.id) ? messages : [...messages, msg]
+    );
+
+  const insertMessages = (msgs: Message[]) =>
+    setMessagesState((messages) => uniqueMessages([...msgs, ...messages]));
+
+  const deleteMessage = (id: number) =>
+    setMessagesState((messages) => messages.filter((m) => m.id !== id));
+
+  const editMessage = (id: number, contents: string) =>
+    setMessagesState((messages) =>
+      messages.map((m) => (m.id === id ? { ...m, contents } : m))
+    );
+
+  const clearMessages = () => setMessagesState([]);
 
   return {
-    messages: messages,
+    messages,
-    setMessages: (msgs) => setMessages(msgs),
+    setMessages,
-    addMessage: (msg) =>
+    addMessage,
-      setMessages((messages) =>
+    deleteMessage,
-        messages.some((m) => m.id === msg.id) ? messages : [...messages, msg]
+    editMessage,
-      ),
+    clearMessages,
-    deleteMessage: (id: number) =>
+    insertMessages,
-      setMessages((messages) => messages.filter((m) => m.id !== id)),
-    editMessage: (id: number, contents) =>
-      setMessages((messages) =>
-        messages.map((m) => (m.id === id ? { ...m, contents } : m))
-      ),
-    clearMessages: () => setMessages([]),
-    insertMessages: (msgs) => setMessages((m) => [...msgs, ...m]),
   };
 }

src/hooks/use-node.ts

@@ -0,0 +1,47 @@
+import { RefObject, useEffect, useRef } from "react";
+import { MessageStore } from "./use-messages";
+import { NodeCallbacks, NodeClient } from "@/lib/node";
+import { asyncWrapProviders } from "async_hooks";
+
+export default function useNode(
+  nodeId: string,
+  messageStore: MessageStore,
+  callbacks: NodeCallbacks
+) {
+  const wsClientRef = useRef<NodeClient | null>(null);
+  useEffect(() => getNode(nodeId, messageStore, callbacks, wsClientRef), []);
+  return wsClientRef;
+}
+
+export function useDynNode(
+  nodeId: string | undefined | null,
+  messageStore: MessageStore,
+  callbacks: NodeCallbacks
+) {
+  const wsClientRef = useRef<NodeClient | null>(null);
+
+  useEffect(() => {
+    if (!nodeId) return;
+    return getNode(nodeId, messageStore, callbacks, wsClientRef);
+  }, [nodeId]);
+
+  return wsClientRef;
+}
+
+export function getNode(
+  nodeId: string,
+  messageStore: MessageStore,
+  callbacks: NodeCallbacks,
+  ref: RefObject<NodeClient | null>
+) {
+  const client = new NodeClient(nodeId, messageStore, callbacks);
+
+  ref.current = client;
+  client.connect();
+
+  return () => client.disconnect();
+}
+
+export function useNullNode() {
+  return useRef<NodeClient | null>(null);
+}

src/hooks/use-once.ts

@@ -1,7 +1,7 @@
 import { useEffect, useRef } from "react";
 
 export function useEffectOnceWhenReady(
-  effect: () => void,
+  effect: () => any,
   deps: any[],
   testers: (undefined | ((m: any) => boolean))[]
 ) {
@@ -13,6 +13,7 @@ export function useEffectOnceWhenReady(
     );
     if (!allReady || hasRun.current) return;
     hasRun.current = true;
-    effect();
+
+    return effect();
   }, deps);
 }

src/lib/auth.ts

@@ -1,165 +0,0 @@
-import { Message, Server } from "@/types/types";
-import { Dispatch, RefObject, SetStateAction } from "react";
-import { toast } from "sonner";
-import { MessageStore } from "@/hooks/use-messages";
-import { IndicatorContext, ServerMessage } from "@/types/protocol";
-import { authRelay } from "@/actions/auth";
-import { getServerById } from "@/actions/get-server";
-import { playPCM16 } from "./audio";
-
-type AuthParams = {
-  id: string;
-  wsRef: RefObject<WebSocket | null>;
-  setServer?: Dispatch<SetStateAction<Server | undefined>>;
-  onNewMessage?: (m: Message) => void;
-  messageStore: MessageStore;
-  onIndicator?: (i: IndicatorContext) => void;
-  onVoice?: (userId: string, bytes: number[]) => void;
-  onVoiceJoin?: (userId: string, channelId: string, voiceId: number) => void;
-  onVoiceLeave?: (userId: string, channelId: string, voiceId: number) => void;
-  onSpeaking?: (voiceId: number) => void;
-};
-
-export default async function auth({
-  id,
-  wsRef,
-  setServer,
-  onNewMessage,
-  messageStore,
-  onIndicator,
-  onVoiceJoin,
-  onVoiceLeave,
-  onSpeaking,
-}: AuthParams) {
-  console.log("Authenticating with server id:", id);
-
-  if (!id) return;
-
-  const ip = (await getServerById(id)).address;
-
-  console.log("Authenticating with server at:", ip);
-  const server_auth = await authRelay(id);
-
-  var name = id;
-
-  // Open the WebSocket connection
-  const ws = new WebSocket(`wss://${ip}`);
-  wsRef.current = ws;
-
-  ws.onopen = () => {
-    console.log("Connected to WebSocket:", ip);
-  };
-
-  ws.onmessage = async (m) => {
-    if (m.data instanceof Blob) {
-      const buffer = await m.data.arrayBuffer();
-      const bytes = new Uint8Array(buffer);
-      playPCM16(bytes);
-      if (onSpeaking) onSpeaking(new DataView(buffer).getUint16(0, true));
-      return;
-    }
-
-    var data = JSON.parse(m.data);
-    console.log("Message received:", data);
-
-    if (data.error) {
-      toast.error(`Error from ${name} (${data.error})`, {
-        description: data.message,
-      });
-      return;
-    }
-
-    if (data.version) {
-      if (setServer) setServer({ channels: [], ...data });
-      ws.send(
-        JSON.stringify({
-          version: "0.0.1",
-          auth_token: server_auth,
-        })
-      );
-      name = data.name;
-      return;
-    }
-
-    const msg = data as ServerMessage;
-
-    switch (msg.type) {
-      case "authenticated":
-        if (onIndicator) msg.params.indicators.forEach(onIndicator);
-        if (onVoiceJoin)
-          Object.entries(msg.params.voice_chat).forEach(([channelId, users]) =>
-            Object.entries(users).forEach(([userId, voiceId]) =>
-              onVoiceJoin(userId, channelId, voiceId)
-            )
-          );
-        break;
-
-      case "message_create":
-        if (onNewMessage) onNewMessage(msg.params as Message);
-        messageStore.addMessage(msg.params as Message);
-        break;
-
-      case "message_delete":
-        messageStore.deleteMessage(msg.params.message_id);
-        break;
-
-      case "message_update":
-        messageStore.editMessage(msg.params.message_id, msg.params.contents);
-        break;
-      case "chunk":
-        messageStore.insertMessages(msg.params);
-        break;
-      case "indicator":
-        if (!onIndicator) break;
-        onIndicator(msg.params);
-        break;
-      case "voice_join":
-        if (!onVoiceJoin) break;
-        onVoiceJoin(
-          msg.params.user_id,
-          msg.params.channel_id,
-          msg.params.voice_id
-        );
-        break;
-      case "voice_leave":
-        if (!onVoiceLeave) break;
-        onVoiceLeave(
-          msg.params.user_id,
-          msg.params.channel_id,
-          msg.params.voice_id
-        );
-        break;
-    }
-  };
-
-  ws.onerror = (err) => {
-    console.error("WebSocket error:", err);
-
-    let description = "An unknown WebSocket error occurred.";
-
-    if (err instanceof ErrorEvent) {
-      // Browser WebSocket errors show up here
-      description =
-        err.message || err.error?.message || "WebSocket connection failed.";
-    } else if (err instanceof Event) {
-      // Generic Event (no info)
-      description = "WebSocket encountered a network issue.";
-    } else if (err && typeof err === "object") {
-      // Custom data if something special got passed
-      description = JSON.stringify(err, null, 2);
-    } else {
-      description = String(err);
-    }
-
-    toast.error(`WebSocket error (${name})`, { description });
-  };
-
-  ws.onclose = () => {
-    console.log("WebSocket closed: ", ip);
-  };
-
-  // Cleanup on unmount or id change
-  return () => {
-    ws.close();
-  };
-}

src/lib/node.ts

@@ -0,0 +1,203 @@
+import { Message, Server } from "@/types/types";
+import { toast } from "sonner";
+import { MessageStore } from "@/hooks/use-messages";
+import { IndicatorContext, ServerMessage } from "@/types/protocol";
+import { authRelay } from "@/actions/auth";
+import { getServerById } from "@/actions/get-server";
+import { playPCM16 } from "./audio";
+
+export interface NodeCallbacks {
+  setServer?: (s: Server) => void;
+  onNewMessage?: (m: Message) => void;
+  onIndicator?: (i: IndicatorContext) => void;
+  onVoiceJoin?: (userId: string, channelId: string, voiceId: number) => void;
+  onVoiceLeave?: (userId: string, channelId: string, voiceId: number) => void;
+  onSpeaking?: (voiceId: number) => void;
+}
+
+export class NodeClient {
+  private ws: WebSocket | null = null;
+  private name: string;
+  private readonly serverId: string;
+  private readonly messageStore: MessageStore;
+  private readonly callbacks: NodeCallbacks;
+  private ready: boolean;
+  private readyPromise: Promise<void>;
+  private resolveReady!: () => void;
+
+  constructor(
+    serverId: string,
+    messageStore: MessageStore,
+    callbacks: NodeCallbacks = {}
+  ) {
+    this.serverId = serverId;
+    this.messageStore = messageStore;
+    this.callbacks = callbacks;
+    this.name = serverId;
+    this.ready = false;
+    this.readyPromise = new Promise((resolve) => {
+      this.resolveReady = resolve;
+    });
+  }
+
+  get socket() {
+    return this.ws;
+  }
+
+  async connect() {
+    if (!this.serverId) return;
+
+    const server = await getServerById(this.serverId);
+    const authToken = await authRelay(this.serverId);
+
+    this.ws = new WebSocket(`wss://${server.address}`);
+
+    this.ws.onopen = () => {
+      console.log("WebSocket connected:", server.address);
+    };
+
+    this.ws.onmessage = (e) => this.handleMessage(e, authToken);
+    this.ws.onerror = (e) => this.handleError(e);
+    this.ws.onclose = () => {
+      console.log("WebSocket closed:", server.address);
+    };
+  }
+
+  disconnect() {
+    console.log("Disconnecting");
+    this.ws?.close();
+    this.ws = null;
+  }
+
+  async waitUntilReady() {
+    if (this.ready) return;
+    await this.readyPromise;
+  }
+
+  loadInitialMessages(channelId: string) {
+    this.waitUntilReady().then(() => {
+      this.socket?.send(
+        JSON.stringify({
+          type: "load_chunk",
+          params: {
+            channel_id: channelId,
+            chunk_id: 0,
+          },
+        })
+      );
+    });
+  }
+
+  private async handleMessage(e: MessageEvent, authToken: string) {
+    // Voice packet
+    if (e.data instanceof Blob) {
+      const buffer = await e.data.arrayBuffer();
+      const bytes = new Uint8Array(buffer);
+
+      playPCM16(bytes);
+
+      if (this.callbacks.onSpeaking) {
+        const voiceId = new DataView(buffer).getUint16(0, true);
+        this.callbacks.onSpeaking(voiceId);
+      }
+      return;
+    }
+
+    const data = JSON.parse(e.data);
+
+    console.log("Message received:", data);
+
+    if (data.error) {
+      toast.error(`Error from ${this.name} (${data.error})`, {
+        description: data.message,
+      });
+      return;
+    }
+
+    // Handshake
+    if (data.version) {
+      this.callbacks.setServer?.({ channels: [], ...data });
+
+      this.ws?.send(
+        JSON.stringify({
+          version: "0.0.1",
+          auth_token: authToken,
+        })
+      );
+
+      this.name = data.name;
+      return;
+    }
+
+    this.handleServerMessage(data as ServerMessage);
+  }
+
+  private handleServerMessage(msg: ServerMessage) {
+    switch (msg.type) {
+      case "authenticated":
+        msg.params.indicators.forEach((i) => this.callbacks.onIndicator?.(i));
+
+        Object.entries(msg.params.voice_chat).forEach(([channelId, users]) =>
+          Object.entries(users).forEach(([userId, voiceId]) =>
+            this.callbacks.onVoiceJoin?.(userId, channelId, voiceId)
+          )
+        );
+
+        this.ready = true;
+        this.resolveReady();
+        break;
+
+      case "message_create":
+        this.callbacks.onNewMessage?.(msg.params);
+        this.messageStore.addMessage(msg.params);
+        break;
+
+      case "message_delete":
+        this.messageStore.deleteMessage(msg.params.message_id);
+        break;
+
+      case "message_update":
+        this.messageStore.editMessage(
+          msg.params.message_id,
+          msg.params.contents
+        );
+        break;
+
+      case "chunk":
+        this.messageStore.insertMessages(msg.params);
+        break;
+
+      case "indicator":
+        this.callbacks.onIndicator?.(msg.params);
+        break;
+
+      case "voice_join":
+        this.callbacks.onVoiceJoin?.(
+          msg.params.user_id,
+          msg.params.channel_id,
+          msg.params.voice_id
+        );
+        break;
+
+      case "voice_leave":
+        this.callbacks.onVoiceLeave?.(
+          msg.params.user_id,
+          msg.params.channel_id,
+          msg.params.voice_id
+        );
+        break;
+    }
+  }
+
+  private handleError(err: Event) {
+    console.error("WebSocket error:", err);
+
+    let description = "WebSocket connection failed.";
+
+    if (err instanceof ErrorEvent) {
+      description = err.message || description;
+    }
+
+    toast.error(`WebSocket error (${this.name})`, { description });
+  }
+}

src/types/app.ts

@@ -1,14 +1,15 @@
 import { Dispatch, RefObject, SetStateAction } from "react";
 import { StringMap } from "./typeUtils";
 import { UserProfile } from "@/hooks/get-user";
-import { Channel, Message, Server } from "./types";
+import { Channel, Server } from "./types";
 import { ClientSettings } from "./settings";
 import { MessageStore } from "@/hooks/use-messages";
+import { NodeClient } from "@/lib/node";
 
 export type SetState<T> = Dispatch<SetStateAction<T>>;
 
 export default interface App {
-  node: RefObject<WebSocket | null>;
+  node: RefObject<NodeClient | null>;
   profile: UserProfile | null;
 
   server: Server | undefined;