diff --git a/ui/colony/src/App.tsx b/ui/colony/src/App.tsx index 79a3302..3c74db2 100644 --- a/ui/colony/src/App.tsx +++ b/ui/colony/src/App.tsx @@ -42,13 +42,17 @@ export default function App() { // WebSocket: append new messages in real-time const handleWsMessage = useCallback((msg: Message) => { setMessages((prev) => { - // Deduplicate — the sender also fetches after posting if (prev.some((m) => m.id === msg.id)) return prev; return [...prev, msg]; }); }, []); - useChannelSocket(activeChannelId, handleWsMessage); + // On WS reconnect, refetch full history to catch missed messages + const handleWsReconnect = useCallback(() => { + loadMessages(); + }, [loadMessages]); + + useChannelSocket(activeChannelId, handleWsMessage, handleWsReconnect); useEffect(() => { loadChannels(); }, [loadChannels]); diff --git a/ui/colony/src/hooks/useChannelSocket.ts b/ui/colony/src/hooks/useChannelSocket.ts index df61e8b..0ebdccf 100644 --- a/ui/colony/src/hooks/useChannelSocket.ts +++ b/ui/colony/src/hooks/useChannelSocket.ts @@ -1,6 +1,6 @@ import { useEffect, useRef, useCallback } from "react"; import type { Message } from "@/types/Message"; -import { getCurrentUsername } from "@/api"; +import { getCurrentUsername, getMessages } from "@/api"; interface WsMessageEvent { event: "message"; @@ -16,19 +16,22 @@ type WsEvent = WsMessageEvent | WsConnectedEvent; export function useChannelSocket( channelId: string | null, onMessage: (msg: Message) => void, + onReconnect: () => void, ) { const wsRef = useRef(null); const reconnectTimer = useRef | null>(null); + const intentionalClose = useRef(false); const connect = useCallback(() => { if (!channelId) return; + intentionalClose.current = false; + const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"; const host = window.location.host; const ws = new WebSocket(`${protocol}//${host}/ws/${channelId}`); ws.onopen = () => { - // Send auth message ws.send(JSON.stringify({ type: "auth", user: getCurrentUsername(), @@ -36,23 +39,38 @@ export function useChannelSocket( }; ws.onmessage = (e) => { - const event: WsEvent = JSON.parse(e.data); - if (event.event === "message") { - onMessage(event.data); + try { + const event: WsEvent = JSON.parse(e.data); + if (event.event === "message") { + onMessage(event.data); + } else if (event.event === "connected") { + // Refetch history on reconnect to catch missed messages + onReconnect(); + } + } catch { + // Ignore malformed messages } }; + ws.onerror = () => { + // Will trigger onclose + }; + ws.onclose = () => { - // Reconnect after 3s - reconnectTimer.current = setTimeout(connect, 3000); + // Only reconnect if this wasn't an intentional teardown + if (!intentionalClose.current) { + reconnectTimer.current = setTimeout(connect, 3000); + } }; wsRef.current = ws; - }, [channelId, onMessage]); + }, [channelId, onMessage, onReconnect]); useEffect(() => { connect(); return () => { + // Mark as intentional so onclose doesn't reconnect + intentionalClose.current = true; if (reconnectTimer.current) clearTimeout(reconnectTimer.current); wsRef.current?.close(); wsRef.current = null;