--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/love/epi/src/hooks/useChatWebsocket.ts	Mon Dec 01 20:35:56 2025 -0800
@@ -0,0 +1,177 @@
+import { wsUrl } from '@/utils';
+import { useEffect } from 'react';
+
+export type Payload = {
+  chatId: string;
+  content: string;
+  action: 'append' | 'done';
+} | {
+  chatId: string;
+  title: string;
+  action:'title_updated';
+} | {
+  chatId: string;
+  url: string;
+  action:'image';
+}
+
+export type OnMessage = (payload: Payload) => void;
+
+type WebSocketEvent = {
+  chatId: string;
+  payload: Payload;
+};
+
+// TODO: Make this into class so we can mock for test.
+const wsSingleton = new Map<string, WebSocket>();
+const listeners = new Map<string, Set<OnMessage>>();
+const pendingMessages = new Map<string, Array<{ content: string; resolve: () => void }>>();
+
+function getOrCreateWebSocket(chatId: string): WebSocket {
+  let ws = wsSingleton.get(chatId);
+  
+  if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
+    return ws;
+  }
+
+  ws = new WebSocket(wsUrl(`/chats/${chatId}/ws`));
+
+  ws.onopen = () => {
+    console.log(`[WS] Connected for chat ${chatId}`);
+    const queue = pendingMessages.get(chatId);
+    if (queue) {
+      queue.forEach(({ content }) => {
+        ws!.send(JSON.stringify({ role: 'user', content }));
+      });
+      pendingMessages.delete(chatId);
+    }
+  };
+
+  // ADD IT HERE:
+  ws.onmessage = (event) => {
+    const payload = JSON.parse(event.data);
+    console.log('[WS] Received message:', payload);
+    
+    wsEventTarget.dispatchEvent(
+      new CustomEvent('message', {
+        detail: { chatId, payload },
+      }),
+    );
+  };
+
+  ws.onerror = (error) => {
+    console.error(`[WS] Error for chat ${chatId}:`, error);
+  };
+
+  ws.onclose = () => {
+    console.log(`[WS] Closed for chat ${chatId}`);
+    wsSingleton.delete(chatId);
+  };
+
+  wsSingleton.set(chatId, ws);
+  return ws;
+}
+
+// TODO: This could be not done rather cancel?
+export function broadcastDone(chatId: string) {
+  const set = listeners.get(chatId);
+  // TODO: This does not update the histroy to be canceled on so there is decrepency but it should be fine for now.
+  if (set) {
+    set.forEach((cb) => cb({ chatId, content: '', action: 'done' }));
+  }
+}
+
+export function useChatWebSocket(
+  chatId: string | null,
+  onMessage: OnMessage,
+  onLoadingChange: (loading: boolean) => void,
+) {
+
+  useEffect(() => {
+    if (!chatId) return;
+
+    // Only register listener to given chatId as chatId websocket has not been made yet. 
+    let set = listeners.get(chatId);
+    if (!set) {
+      set = new Set();
+      listeners.set(chatId, set);
+    }
+    set.add(onMessage);
+
+    return () => {
+      if (!chatId) return;
+      const set = listeners.get(chatId);
+      if (set) {
+        set.delete(onMessage);
+        if (set.size === 0) {
+          listeners.delete(chatId);
+          const ws = wsSingleton.get(chatId);
+          if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
+            ws.close();
+          }
+        }
+      }
+    };
+  }, [chatId, onMessage]);
+
+  const sendMessage = (content: string) => {
+    if (!chatId) return;
+
+    const ws = getOrCreateWebSocket(chatId);
+
+    onLoadingChange(true);
+
+    if (ws.readyState === WebSocket.OPEN) {
+      ws.send(JSON.stringify({ role: 'user', content }));
+    } else if (ws.readyState === WebSocket.CONNECTING) {
+      // Queue message until open
+      let queue = pendingMessages.get(chatId);
+      if (!queue) {
+        queue = [];
+        pendingMessages.set(chatId, queue);
+      }
+      console.log('Queue: ', queue);
+      if (queue.length > 1) return;
+      queue.push({ content, resolve: () => {} });
+    } else {
+      console.warn('[WS] WebSocket not in usable state', ws.readyState);
+      onLoadingChange(false);
+    }
+  };
+
+  return { sendMessage };
+}
+
+const wsEventTarget = new EventTarget();
+
+export function subscribeToChat(chatId: string, callback: (payload: Payload) => void) {
+  const handler = (event: Event) => {
+    const { chatId: eventChatId, payload } = (event as CustomEvent<WebSocketEvent>).detail;
+    if (eventChatId === chatId) {
+      callback(payload);
+    }
+  };
+  
+  wsEventTarget.addEventListener('message', handler);
+  
+  return () => wsEventTarget.removeEventListener('message', handler);
+}
+
+export function sendMessageChatId(content: string, chatId: string): void {
+  if (!content.trim() || !chatId) return;
+  
+  const ws = getOrCreateWebSocket(chatId);
+  
+  if (ws.readyState === WebSocket.OPEN) {
+    ws.send(JSON.stringify({ role: 'user', content }));
+  } else if (ws.readyState === WebSocket.CONNECTING) {
+    let queue = pendingMessages.get(chatId);
+    if (!queue) {
+      queue = [];
+      pendingMessages.set(chatId, queue);
+    }
+    if (!queue.some(m => m.content === content)) {
+      queue.push({ content, resolve: () => {} });
+    }
+  }
+}