| 1 | import { useCallback, useEffect, useRef, useState } from "react"; |
| 2 | |
| 3 | import { useClient } from "@/providers/ClientProvider"; |
| 4 | import i18n from "@/i18n"; |
| 5 | import { |
| 6 | ApiError, |
| 7 | deleteSession as apiDeleteSession, |
| 8 | fetchSessionMessages, |
| 9 | listSessions, |
| 10 | } from "@/lib/api"; |
| 11 | import { deriveTitle } from "@/lib/format"; |
| 12 | import { wireSessionMessages } from "@/lib/questions"; |
| 13 | import { webuiSessionKey } from "@/lib/session-key"; |
| 14 | import type { ChatSummary, UIMessage } from "@/lib/types"; |
| 15 | |
| 16 | const EMPTY_MESSAGES: UIMessage[] = []; |
| 17 | |
| 18 | /** Merge server rows with local-only chats that are not persisted yet. */ |
| 19 | export function mergeSessionLists( |
| 20 | prev: ChatSummary[], |
| 21 | serverRows: ChatSummary[], |
| 22 | localKeys: Set<string>, |
| 23 | ): ChatSummary[] { |
| 24 | const serverByKey = new Map(serverRows.map((row) => [row.key, row])); |
| 25 | |
| 26 | for (const key of serverByKey.keys()) { |
| 27 | localKeys.delete(key); |
| 28 | } |
| 29 | |
| 30 | const locals = prev.filter( |
| 31 | (session) => |
| 32 | localKeys.has(session.key) && !serverByKey.has(session.key), |
| 33 | ); |
| 34 | |
| 35 | const sortedServer = [...serverRows].sort( |
| 36 | (a, b) => |
| 37 | Date.parse(b.updatedAt ?? "") - Date.parse(a.updatedAt ?? ""), |
| 38 | ); |
| 39 | |
| 40 | return [...locals, ...sortedServer]; |
| 41 | } |
| 42 | |
| 43 | /** Sidebar state: fetches the full session list and exposes create / delete actions. */ |
| 44 | export function useSessions(): { |
| 45 | sessions: ChatSummary[]; |
| 46 | loading: boolean; |
| 47 | error: string | null; |
| 48 | refresh: () => Promise<void>; |
| 49 | createChat: (options?: { autoGenerate?: boolean }) => Promise<string>; |
| 50 | deleteChat: (key: string) => Promise<void>; |
| 51 | } { |
| 52 | const { client, token, webuiUserId } = useClient(); |
| 53 | const [sessions, setSessions] = useState<ChatSummary[]>([]); |
| 54 | const [loading, setLoading] = useState(true); |
| 55 | const [error, setError] = useState<string | null>(null); |
| 56 | const tokenRef = useRef(token); |
| 57 | tokenRef.current = token; |
| 58 | const localSessionKeysRef = useRef(new Set<string>()); |
| 59 | |
| 60 | const refresh = useCallback(async () => { |
| 61 | try { |
| 62 | setLoading(true); |
| 63 | const rows = await listSessions(tokenRef.current); |
| 64 | setSessions((prev) => |
| 65 | mergeSessionLists(prev, rows, localSessionKeysRef.current), |
| 66 | ); |
| 67 | setError(null); |
| 68 | } catch (e) { |
| 69 | const msg = |
| 70 | e instanceof ApiError ? `HTTP ${e.status}` : (e as Error).message; |
| 71 | setError(msg); |
| 72 | } finally { |
| 73 | setLoading(false); |
| 74 | } |
| 75 | }, []); |
| 76 | |
| 77 | useEffect(() => { |
| 78 | void refresh(); |
| 79 | }, [refresh]); |
| 80 | |
| 81 | const createChat = useCallback( |
| 82 | async (options?: { autoGenerate?: boolean }): Promise<string> => { |
| 83 | const chatId = await client.newChat(5_000, options); |
| 84 | const key = webuiSessionKey(webuiUserId, chatId); |
| 85 | localSessionKeysRef.current.add(key); |
| 86 | setSessions((prev) => [ |
| 87 | { |
| 88 | key, |
| 89 | channel: "websocket", |
| 90 | chatId, |
| 91 | createdAt: new Date().toISOString(), |
| 92 | updatedAt: new Date().toISOString(), |
| 93 | preview: "", |
| 94 | source: "stepwise", |
| 95 | autoGenerate: Boolean(options?.autoGenerate), |
| 96 | }, |
| 97 | ...prev.filter((s) => s.key !== key), |
| 98 | ]); |
| 99 | return chatId; |
| 100 | }, |
| 101 | [client, webuiUserId], |
| 102 | ); |
| 103 | |
| 104 | const deleteChat = useCallback(async (key: string) => { |
| 105 | await apiDeleteSession(tokenRef.current, key); |
| 106 | localSessionKeysRef.current.delete(key); |
| 107 | setSessions((prev) => prev.filter((s) => s.key !== key)); |
| 108 | }, []); |
| 109 | |
| 110 | return { sessions, loading, error, refresh, createChat, deleteChat }; |
| 111 | } |
| 112 | |
| 113 | /** Lazy-load a session's on-disk messages the first time the UI displays it. */ |
| 114 | export function useSessionHistory(key: string | null): { |
| 115 | messages: UIMessage[]; |
| 116 | loading: boolean; |
| 117 | error: string | null; |
| 118 | } { |
| 119 | const { token } = useClient(); |
| 120 | const [state, setState] = useState<{ |
| 121 | key: string | null; |
| 122 | messages: UIMessage[]; |
| 123 | loading: boolean; |
| 124 | error: string | null; |
| 125 | }>({ |
| 126 | key: null, |
| 127 | messages: [], |
| 128 | loading: false, |
| 129 | error: null, |
| 130 | }); |
| 131 | |
| 132 | useEffect(() => { |
| 133 | if (!key) { |
| 134 | setState({ |
| 135 | key: null, |
| 136 | messages: [], |
| 137 | loading: false, |
| 138 | error: null, |
| 139 | }); |
| 140 | return; |
| 141 | } |
| 142 | let cancelled = false; |
| 143 | // Mark the new key as loading immediately so callers never see stale |
| 144 | // messages from the previous session during the render right after a switch. |
| 145 | setState({ |
| 146 | key, |
| 147 | messages: [], |
| 148 | loading: true, |
| 149 | error: null, |
| 150 | }); |
| 151 | (async () => { |
| 152 | try { |
| 153 | const body = await fetchSessionMessages(token, key); |
| 154 | if (cancelled) return; |
| 155 | const ui: UIMessage[] = wireSessionMessages(body.messages); |
| 156 | setState({ |
| 157 | key, |
| 158 | messages: ui, |
| 159 | loading: false, |
| 160 | error: null, |
| 161 | }); |
| 162 | } catch (e) { |
| 163 | if (cancelled) return; |
| 164 | // A 404 just means the session hasn't been persisted yet (brand-new |
| 165 | // chat, first message not sent). That's a normal state, not an error. |
| 166 | if (e instanceof ApiError && e.status === 404) { |
| 167 | setState({ |
| 168 | key, |
| 169 | messages: [], |
| 170 | loading: false, |
| 171 | error: null, |
| 172 | }); |
| 173 | } else { |
| 174 | setState({ |
| 175 | key, |
| 176 | messages: [], |
| 177 | loading: false, |
| 178 | error: (e as Error).message, |
| 179 | }); |
| 180 | } |
| 181 | } |
| 182 | })(); |
| 183 | return () => { |
| 184 | cancelled = true; |
| 185 | }; |
| 186 | }, [key, token]); |
| 187 | |
| 188 | if (!key) { |
| 189 | return { messages: EMPTY_MESSAGES, loading: false, error: null }; |
| 190 | } |
| 191 | |
| 192 | // Even before the effect above commits its loading state, never surface the |
| 193 | // previous session's payload for a brand-new key. |
| 194 | if (state.key !== key) { |
| 195 | return { messages: EMPTY_MESSAGES, loading: true, error: null }; |
| 196 | } |
| 197 | |
| 198 | return { |
| 199 | messages: state.messages, |
| 200 | loading: state.loading, |
| 201 | error: state.error, |
| 202 | }; |
| 203 | } |
| 204 | |
| 205 | /** Produce a compact display title for a session. */ |
| 206 | export function sessionTitle( |
| 207 | session: ChatSummary, |
| 208 | firstUserMessage?: string, |
| 209 | ): string { |
| 210 | return deriveTitle( |
| 211 | firstUserMessage || session.preview, |
| 212 | i18n.t("chat.newChat"), |
| 213 | ); |
| 214 | } |
| 215 |