diff --git a/src/client.ts b/src/client.ts
index b3bcc74..b684f80 100644
--- a/src/client.ts
+++ b/src/client.ts
@@ -58,9 +58,16 @@ export function createClient(config: {
     token,
   };
 
-  const socket = RoomsSocket({
-    config: socketConfig,
-  });
+  let socket: ReturnType<typeof RoomsSocket> | null = null;
+
+  const getSocket = () => {
+    if (!socket) {
+      socket = RoomsSocket({
+        config: socketConfig,
+      });
+    }
+    return socket;
+  };
 
   const headers = {
     ...optionalHeaders,
@@ -113,14 +120,16 @@ export function createClient(config: {
     functions: createFunctionsModule(functionsAxiosClient, appId),
     agents: createAgentsModule({
       axios: axiosClient,
-      socket,
+      getSocket,
       appId,
       serverUrl,
       token,
     }),
     appLogs: createAppLogsModule(axiosClient, appId),
     cleanup: () => {
-      socket.disconnect();
+      if (socket) {
+        socket.disconnect();
+      }
     },
   };
 
@@ -132,14 +141,16 @@ export function createClient(config: {
     functions: createFunctionsModule(serviceRoleFunctionsAxiosClient, appId),
     agents: createAgentsModule({
       axios: serviceRoleAxiosClient,
-      socket,
+      getSocket,
       appId,
       serverUrl,
       token,
     }),
     appLogs: createAppLogsModule(serviceRoleAxiosClient, appId),
     cleanup: () => {
-      socket.disconnect();
+      if (socket) {
+        socket.disconnect();
+      }
     },
   };
 
@@ -178,9 +189,12 @@ export function createClient(config: {
      */
     setToken(newToken: string) {
       userModules.auth.setToken(newToken);
-      socket.updateConfig({
-        token: newToken,
-      });
+      if (socket) {
+        socket.updateConfig({
+          token: newToken,
+        });
+      }
+      socketConfig.token = newToken;
     },
 
     /**
diff --git a/src/modules/agents.ts b/src/modules/agents.ts
index e1cca5c..c573cd1 100644
--- a/src/modules/agents.ts
+++ b/src/modules/agents.ts
@@ -6,7 +6,7 @@ import { getAccessToken } from "../utils/auth-utils.js";
 
 export type AgentsModuleConfig = {
   axios: AxiosInstance;
-  socket: ReturnType<typeof RoomsSocket>;
+  getSocket: () => ReturnType<typeof RoomsSocket>;
   appId: string;
   serverUrl?: string;
   token?: string;
@@ -14,7 +14,7 @@ export type AgentsModuleConfig = {
 
 export function createAgentsModule({
   axios,
-  socket,
+  getSocket,
   appId,
   serverUrl,
   token,
@@ -52,6 +52,7 @@ export function createAgentsModule({
     message: AgentMessage
   ) => {
     const room = `/agent-conversations/${conversation.id}`;
+    const socket = getSocket();
     await socket.updateModel(
       room,
       {
@@ -70,6 +71,7 @@ export function createAgentsModule({
     onUpdate?: (conversation: AgentConversation) => void
   ) => {
     const room = `/agent-conversations/${conversationId}`;
+    const socket = getSocket();
     return socket.subscribeToRoom(room, {
       connect: () => {},
       update_model: ({ data: jsonStr }) => {