Mercurial

comparison love/epi/src/hooks/useChatWebsocket.ts @ 38:cf9caa4abc3e

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
[Love] FE and BE. Can chat and render images. Also created MCP for powerpoint generations.
author MrJuneJune <me@mrjunejune.com>
date Mon, 01 Dec 2025 20:35:56 -0800
parents
children
comparison
equal deleted inserted replaced
37:fb9bcd3145cb 38:cf9caa4abc3e
1 import { wsUrl } from '@/utils';
2 import { useEffect } from 'react';
3
4 export type Payload = {
5 chatId: string;
6 content: string;
7 action: 'append' | 'done';
8 } | {
9 chatId: string;
10 title: string;
11 action:'title_updated';
12 } | {
13 chatId: string;
14 url: string;
15 action:'image';
16 }
17
18 export type OnMessage = (payload: Payload) => void;
19
20 type WebSocketEvent = {
21 chatId: string;
22 payload: Payload;
23 };
24
25 // TODO: Make this into class so we can mock for test.
26 const wsSingleton = new Map<string, WebSocket>();
27 const listeners = new Map<string, Set<OnMessage>>();
28 const pendingMessages = new Map<string, Array<{ content: string; resolve: () => void }>>();
29
30 function getOrCreateWebSocket(chatId: string): WebSocket {
31 let ws = wsSingleton.get(chatId);
32
33 if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
34 return ws;
35 }
36
37 ws = new WebSocket(wsUrl(`/chats/${chatId}/ws`));
38
39 ws.onopen = () => {
40 console.log(`[WS] Connected for chat ${chatId}`);
41 const queue = pendingMessages.get(chatId);
42 if (queue) {
43 queue.forEach(({ content }) => {
44 ws!.send(JSON.stringify({ role: 'user', content }));
45 });
46 pendingMessages.delete(chatId);
47 }
48 };
49
50 // ADD IT HERE:
51 ws.onmessage = (event) => {
52 const payload = JSON.parse(event.data);
53 console.log('[WS] Received message:', payload);
54
55 wsEventTarget.dispatchEvent(
56 new CustomEvent('message', {
57 detail: { chatId, payload },
58 }),
59 );
60 };
61
62 ws.onerror = (error) => {
63 console.error(`[WS] Error for chat ${chatId}:`, error);
64 };
65
66 ws.onclose = () => {
67 console.log(`[WS] Closed for chat ${chatId}`);
68 wsSingleton.delete(chatId);
69 };
70
71 wsSingleton.set(chatId, ws);
72 return ws;
73 }
74
75 // TODO: This could be not done rather cancel?
76 export function broadcastDone(chatId: string) {
77 const set = listeners.get(chatId);
78 // TODO: This does not update the histroy to be canceled on so there is decrepency but it should be fine for now.
79 if (set) {
80 set.forEach((cb) => cb({ chatId, content: '', action: 'done' }));
81 }
82 }
83
84 export function useChatWebSocket(
85 chatId: string | null,
86 onMessage: OnMessage,
87 onLoadingChange: (loading: boolean) => void,
88 ) {
89
90 useEffect(() => {
91 if (!chatId) return;
92
93 // Only register listener to given chatId as chatId websocket has not been made yet.
94 let set = listeners.get(chatId);
95 if (!set) {
96 set = new Set();
97 listeners.set(chatId, set);
98 }
99 set.add(onMessage);
100
101 return () => {
102 if (!chatId) return;
103 const set = listeners.get(chatId);
104 if (set) {
105 set.delete(onMessage);
106 if (set.size === 0) {
107 listeners.delete(chatId);
108 const ws = wsSingleton.get(chatId);
109 if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
110 ws.close();
111 }
112 }
113 }
114 };
115 }, [chatId, onMessage]);
116
117 const sendMessage = (content: string) => {
118 if (!chatId) return;
119
120 const ws = getOrCreateWebSocket(chatId);
121
122 onLoadingChange(true);
123
124 if (ws.readyState === WebSocket.OPEN) {
125 ws.send(JSON.stringify({ role: 'user', content }));
126 } else if (ws.readyState === WebSocket.CONNECTING) {
127 // Queue message until open
128 let queue = pendingMessages.get(chatId);
129 if (!queue) {
130 queue = [];
131 pendingMessages.set(chatId, queue);
132 }
133 console.log('Queue: ', queue);
134 if (queue.length > 1) return;
135 queue.push({ content, resolve: () => {} });
136 } else {
137 console.warn('[WS] WebSocket not in usable state', ws.readyState);
138 onLoadingChange(false);
139 }
140 };
141
142 return { sendMessage };
143 }
144
145 const wsEventTarget = new EventTarget();
146
147 export function subscribeToChat(chatId: string, callback: (payload: Payload) => void) {
148 const handler = (event: Event) => {
149 const { chatId: eventChatId, payload } = (event as CustomEvent<WebSocketEvent>).detail;
150 if (eventChatId === chatId) {
151 callback(payload);
152 }
153 };
154
155 wsEventTarget.addEventListener('message', handler);
156
157 return () => wsEventTarget.removeEventListener('message', handler);
158 }
159
160 export function sendMessageChatId(content: string, chatId: string): void {
161 if (!content.trim() || !chatId) return;
162
163 const ws = getOrCreateWebSocket(chatId);
164
165 if (ws.readyState === WebSocket.OPEN) {
166 ws.send(JSON.stringify({ role: 'user', content }));
167 } else if (ws.readyState === WebSocket.CONNECTING) {
168 let queue = pendingMessages.get(chatId);
169 if (!queue) {
170 queue = [];
171 pendingMessages.set(chatId, queue);
172 }
173 if (!queue.some(m => m.content === content)) {
174 queue.push({ content, resolve: () => {} });
175 }
176 }
177 }