| 1 | import { useEffect, useMemo, useRef } from "react"; |
| 2 | |
| 3 | import { useNanobotStream } from "@/hooks/useNanobotStream"; |
| 4 | import { useSessionHistory } from "@/hooks/useSessions"; |
| 5 | import type { ChatSummary, UIMessage } from "@/lib/types"; |
| 6 | |
| 7 | export type ThreadStreamControl = ReturnType<typeof useThreadStream>; |
| 8 | |
| 9 | export function useThreadStream( |
| 10 | session: ChatSummary | null, |
| 11 | options?: { onReplyEnd?: () => void }, |
| 12 | ) { |
| 13 | const chatId = session?.chatId ?? null; |
| 14 | const historyKey = session?.key ?? null; |
| 15 | const { messages: historical, loading } = useSessionHistory(historyKey); |
| 16 | const messageCacheRef = useRef<Map<string, UIMessage[]>>(new Map()); |
| 17 | /** Skip one cache write after chatId changes (messages may still be stale). */ |
| 18 | const skipCacheWriteRef = useRef(false); |
| 19 | |
| 20 | const initial = useMemo(() => { |
| 21 | if (!chatId) return historical; |
| 22 | const cached = messageCacheRef.current.get(chatId); |
| 23 | return cached?.length ? cached : historical; |
| 24 | }, [chatId, historical]); |
| 25 | |
| 26 | const { |
| 27 | messages, |
| 28 | isStreaming, |
| 29 | turnComplete, |
| 30 | send, |
| 31 | answerQuestion, |
| 32 | setMessages, |
| 33 | streamError, |
| 34 | dismissStreamError, |
| 35 | } = useNanobotStream(chatId, initial, { |
| 36 | onReplyEnd: options?.onReplyEnd, |
| 37 | }); |
| 38 | |
| 39 | useEffect(() => { |
| 40 | if (!chatId || loading) return; |
| 41 | const cached = messageCacheRef.current.get(chatId); |
| 42 | // When the user switches away and back, keep the local in-memory thread |
| 43 | // state (including not-yet-persisted messages) instead of replacing it with |
| 44 | // whatever the history endpoint currently knows about. |
| 45 | setMessages(cached && cached.length > 0 ? cached : historical); |
| 46 | // eslint-disable-next-line react-hooks/exhaustive-deps |
| 47 | }, [loading, chatId, historical]); |
| 48 | |
| 49 | useEffect(() => { |
| 50 | if (chatId) return; |
| 51 | setMessages(historical); |
| 52 | }, [chatId, historical, setMessages]); |
| 53 | |
| 54 | useEffect(() => { |
| 55 | skipCacheWriteRef.current = true; |
| 56 | }, [chatId]); |
| 57 | |
| 58 | useEffect(() => { |
| 59 | if (!chatId) return; |
| 60 | if (skipCacheWriteRef.current) { |
| 61 | skipCacheWriteRef.current = false; |
| 62 | return; |
| 63 | } |
| 64 | messageCacheRef.current.set(chatId, messages); |
| 65 | }, [chatId, messages]); |
| 66 | |
| 67 | return { |
| 68 | chatId, |
| 69 | loading, |
| 70 | messages, |
| 71 | isStreaming, |
| 72 | turnComplete, |
| 73 | send, |
| 74 | answerQuestion, |
| 75 | setMessages, |
| 76 | streamError, |
| 77 | dismissStreamError, |
| 78 | }; |
| 79 | } |
| 80 |