"use client"; import { useCallback, useEffect, useRef, useState } from "react"; import { authFetch } from "@/lib/auth/client"; import type { QuotaSnapshot } from "@/lib/chat/quota"; import { MAX_CHAT_MESSAGE_CHARACTERS } from "@/lib/chat/request-limits"; import type { Attachment, ChatMessage, Conversation, ModelId } from "@/lib/chat/types"; const DEFAULT_MODEL_ID: ModelId = "auto"; const GENERIC_ERROR = "Не удалось получить ответ. Попробуйте снова."; const CLIENT_CONVERSATION_PREFIX = "client-conversation-"; const CLIENT_MESSAGE_PREFIX = "client-message-"; class ChatResponseError extends Error {} type OptimisticTurn = { conversationId: string; userMessageId: string; assistantMessageId: string; createdAt: string; clientConversation: boolean; }; export type QuotaState = | { status: "loading" } | { status: "ready"; snapshot: QuotaSnapshot } | { status: "unavailable" }; export type UseChatResult = { conversations: Conversation[]; activeConversation: Conversation | null; activeModelId: ModelId; isGenerating: boolean; canRetryLast: boolean; quota: QuotaState; selectConversation(id: string): void; newConversation(): void; deleteConversation(id: string): void; setModel(id: ModelId): void; sendMessage(content: string, attachments?: Attachment[]): Promise; stopGenerating(): void; retryLast(): Promise; }; export function useChat( userId: string, onUnauthorized?: () => void, initialModelId?: ModelId, ): UseChatResult { const [conversations, setConversations] = useState([]); const [activeConversationId, setActiveConversationId] = useState(null); const [activeModelId, setActiveModelId] = useState(initialModelId ?? DEFAULT_MODEL_ID); const [isGenerating, setIsGenerating] = useState(false); const [quotaState, setQuotaState] = useState<{ userId: string; model: ModelId; value: QuotaState; }>({ userId, model: initialModelId ?? DEFAULT_MODEL_ID, value: { status: "loading" }, }); const conversationsRef = useRef(conversations); const activeConversationIdRef = useRef(activeConversationId); const activeModelIdRef = useRef(activeModelId); const isGeneratingRef = useRef(false); const hasLocalChangesRef = useRef(false); const abortControllerRef = useRef(null); const streamingMessageRef = useRef<{ conversationId: string; messageId: string } | null>(null); const onUnauthorizedRef = useRef(onUnauthorized); const quotaControllerRef = useRef(null); const quotaGenerationRef = useRef(0); const quotaExhaustedRef = useRef(false); const modelPreferenceGenerationRef = useRef(0); const modelPreferenceChainRef = useRef>(Promise.resolve()); const modelPreferenceControllerRef = useRef(null); const modelPreferenceEpochRef = useRef(0); const modelPreferenceMountedRef = useRef(false); useEffect(() => { onUnauthorizedRef.current = onUnauthorized; }, [onUnauthorized]); useEffect(() => { const epoch = modelPreferenceEpochRef.current + 1; modelPreferenceEpochRef.current = epoch; modelPreferenceMountedRef.current = true; modelPreferenceChainRef.current = Promise.resolve(); return () => { if (modelPreferenceEpochRef.current !== epoch) return; modelPreferenceMountedRef.current = false; modelPreferenceEpochRef.current += 1; modelPreferenceGenerationRef.current += 1; modelPreferenceControllerRef.current?.abort(); modelPreferenceControllerRef.current = null; modelPreferenceChainRef.current = Promise.resolve(); }; }, [userId]); const refreshQuota = useCallback(async (model: ModelId, showLoading = true) => { if (model !== activeModelIdRef.current) return; const generation = quotaGenerationRef.current + 1; quotaGenerationRef.current = generation; quotaControllerRef.current?.abort(); const controller = new AbortController(); quotaControllerRef.current = controller; if (showLoading) { quotaExhaustedRef.current = false; setQuotaState({ userId, model, value: { status: "loading" } }); } try { const response = await authFetch( `/api/me/quota?model=${encodeURIComponent(model)}`, { cache: "no-store", signal: controller.signal }, ); if (response.status === 401) { onUnauthorizedRef.current?.(); throw new Error("unauthorized"); } if (!response.ok) throw new Error("quota unavailable"); const snapshot = parseQuotaResponse(await response.json(), model); if (!snapshot) throw new Error("invalid quota"); if ( generation !== quotaGenerationRef.current || controller.signal.aborted || model !== activeModelIdRef.current ) return; quotaExhaustedRef.current = snapshot.exhausted; setQuotaState({ userId, model, value: { status: "ready", snapshot } }); } catch { if ( generation !== quotaGenerationRef.current || controller.signal.aborted || model !== activeModelIdRef.current ) return; quotaExhaustedRef.current = false; setQuotaState({ userId, model, value: { status: "unavailable" } }); } }, [userId]); useEffect(() => { quotaExhaustedRef.current = false; void refreshQuota(activeModelIdRef.current, false); return () => { quotaGenerationRef.current += 1; quotaControllerRef.current?.abort(); }; }, [refreshQuota]); useEffect(() => { const controller = new AbortController(); void authFetch("/api/conversations", { signal: controller.signal }) .then(async (response) => { if (response.status === 401) { onUnauthorizedRef.current?.(); return; } if (!response.ok) return; const body = (await response.json()) as { conversations?: Conversation[] }; const next = Array.isArray(body.conversations) ? body.conversations : []; if (hasLocalChangesRef.current) return; conversationsRef.current = next; setConversations(next); const first = next[0] ?? null; activeConversationIdRef.current = first?.id ?? null; setActiveConversationId(first?.id ?? null); if (first) { activeModelIdRef.current = first.modelId; setActiveModelId(first.modelId); void refreshQuota(first.modelId); } }) .catch(() => undefined); return () => controller.abort(); }, [refreshQuota, userId]); useEffect(() => () => abortControllerRef.current?.abort(), []); const commitConversations = useCallback( (update: (current: Conversation[]) => Conversation[]) => { hasLocalChangesRef.current = true; const next = update(conversationsRef.current); conversationsRef.current = next; setConversations(next); return next; }, [], ); const updateMessage = useCallback( (conversationId: string, messageId: string, update: (message: ChatMessage) => ChatMessage) => { commitConversations((current) => current.map((conversation) => conversation.id === conversationId ? { ...conversation, updatedAt: new Date().toISOString(), messages: conversation.messages.map((message) => message.id === messageId ? update(message) : message, ), } : conversation, ), ); }, [commitConversations], ); const requestAssistant = useCallback( async ( payload: { conversationId?: string; model: ModelId; content?: string; retry?: boolean; attachmentIds?: string[]; }, attachments: Attachment[] = [], ) => { if (isGeneratingRef.current || quotaExhaustedRef.current) return; const optimistic = payload.retry ? null : createOptimisticTurn(payload.conversationId); if (optimistic) { const userMessage: ChatMessage = { id: optimistic.userMessageId, role: "user", content: payload.content ?? "", createdAt: optimistic.createdAt, status: "complete", ...(attachments.length ? { attachments: attachments.map((attachment) => ({ ...attachment })) } : {}), }; const assistantMessage: ChatMessage = { id: optimistic.assistantMessageId, role: "assistant", content: "", createdAt: optimistic.createdAt, status: "streaming", }; commitConversations((current) => { const existing = current.find( (conversation) => conversation.id === optimistic.conversationId, ); if (existing) { return moveFirst( current.map((conversation) => conversation.id === optimistic.conversationId ? { ...conversation, modelId: payload.model, updatedAt: optimistic.createdAt, messages: [ ...conversation.messages, userMessage, assistantMessage, ], } : conversation, ), optimistic.conversationId, ); } return [ { id: optimistic.conversationId, title: createTitle(payload.content ?? "Новый чат"), modelId: payload.model, updatedAt: optimistic.createdAt, messages: [userMessage, assistantMessage], }, ...current, ]; }); if (optimistic.clientConversation) { activeConversationIdRef.current = optimistic.conversationId; setActiveConversationId(optimistic.conversationId); } } const controller = new AbortController(); abortControllerRef.current = controller; isGeneratingRef.current = true; setIsGenerating(true); let responseMessageCommitted = false; try { const response = await authFetch("/api/chat", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(payload), signal: controller.signal, }); if (response.status === 401) { onUnauthorizedRef.current?.(); return; } const conversationId = response.headers.get("X-Conversation-Id"); const assistantId = response.headers.get("X-Assistant-Message-Id"); const userMessageId = response.headers.get("X-User-Message-Id"); if (!conversationId || !assistantId || (optimistic && !userMessageId)) { throw new ChatResponseError(GENERIC_ERROR); } const createdAt = new Date().toISOString(); const commitResponseMessage = (assistantMessage: ChatMessage) => commitConversations((current) => { const sourceConversationId = optimistic?.conversationId ?? conversationId; const existing = current.find( (conversation) => conversation.id === sourceConversationId || conversation.id === conversationId, ); if (existing) { const reconciled = current.map((conversation) => { if ( conversation.id !== sourceConversationId && conversation.id !== conversationId ) return conversation; return { ...conversation, id: conversationId, modelId: payload.model, updatedAt: createdAt, messages: optimistic ? conversation.messages.map((message) => { if (message.id === optimistic.userMessageId) { return { ...message, id: userMessageId as string }; } if (message.id === optimistic.assistantMessageId) { return assistantMessage; } return message; }) : [...conversation.messages, assistantMessage], }; }); return moveFirst(reconciled, conversationId); } return [ { id: conversationId, title: createTitle(payload.content ?? "Новый чат"), modelId: payload.model, updatedAt: createdAt, messages: [ ...(payload.content && userMessageId ? [{ id: userMessageId, role: "user" as const, content: payload.content, createdAt, status: "complete" as const }] : []), assistantMessage, ], }, ...current, ]; }); const activateResponseConversation = activeConversationIdRef.current === optimistic?.conversationId || (!optimistic && activeConversationIdRef.current === payload.conversationId); if (!response.ok || !response.body) { const message = await readSafeError(response); commitResponseMessage({ id: assistantId, role: "assistant", content: message, createdAt, status: "error", }); responseMessageCommitted = true; if (activateResponseConversation) { activeConversationIdRef.current = conversationId; setActiveConversationId(conversationId); } throw new ChatResponseError(message); } const assistantMessage: ChatMessage = { id: assistantId, role: "assistant", content: "", createdAt, status: "streaming", }; commitResponseMessage(assistantMessage); responseMessageCommitted = true; if (activateResponseConversation) { activeConversationIdRef.current = conversationId; setActiveConversationId(conversationId); } streamingMessageRef.current = { conversationId, messageId: assistantId }; const reader = response.body.getReader(); const decoder = new TextDecoder(); while (true) { const { done, value } = await reader.read(); if (done) break; const chunk = decoder.decode(value, { stream: true }); if (chunk) updateMessage(conversationId, assistantId, (message) => ({ ...message, content: message.content + chunk })); } const finalChunk = decoder.decode(); updateMessage(conversationId, assistantId, (message) => ({ ...message, content: message.content + finalChunk, status: "complete" })); } catch (error) { const streaming = streamingMessageRef.current; if (streaming) { updateMessage(streaming.conversationId, streaming.messageId, (message) => ({ ...message, content: message.content || (controller.signal.aborted ? "" : error instanceof ChatResponseError ? error.message : GENERIC_ERROR), status: controller.signal.aborted ? "stopped" : "error", })); } else if (optimistic && !responseMessageCommitted) { updateMessage( optimistic.conversationId, optimistic.assistantMessageId, (message) => ({ ...message, content: controller.signal.aborted ? "" : error instanceof ChatResponseError ? error.message : GENERIC_ERROR, status: controller.signal.aborted ? "stopped" : "error", }), ); } } finally { if (abortControllerRef.current === controller) { abortControllerRef.current = null; streamingMessageRef.current = null; isGeneratingRef.current = false; setIsGenerating(false); void refreshQuota(payload.model); } } }, [commitConversations, refreshQuota, updateMessage], ); const selectConversation = useCallback((id: string) => { const conversation = conversationsRef.current.find((item) => item.id === id); if (!conversation) return; activeConversationIdRef.current = id; activeModelIdRef.current = conversation.modelId; setActiveConversationId(id); setActiveModelId(conversation.modelId); void refreshQuota(conversation.modelId); }, [refreshQuota]); const newConversation = useCallback(() => { activeConversationIdRef.current = null; setActiveConversationId(null); }, []); const deleteConversation = useCallback((id: string) => { const next = commitConversations((current) => current.filter((conversation) => conversation.id !== id)); if (activeConversationIdRef.current === id) { const replacement = next[0] ?? null; activeConversationIdRef.current = replacement?.id ?? null; setActiveConversationId(replacement?.id ?? null); if (replacement) { activeModelIdRef.current = replacement.modelId; setActiveModelId(replacement.modelId); } void refreshQuota(activeModelIdRef.current); } void authFetch(`/api/conversations/${id}`, { method: "DELETE" }).then((response) => { if (response.status === 401) onUnauthorizedRef.current?.(); }); }, [commitConversations, refreshQuota]); const setModel = useCallback((id: ModelId) => { const preferenceEpoch = modelPreferenceEpochRef.current; const preferenceGeneration = modelPreferenceGenerationRef.current + 1; modelPreferenceGenerationRef.current = preferenceGeneration; activeModelIdRef.current = id; setActiveModelId(id); quotaGenerationRef.current += 1; quotaControllerRef.current?.abort(); quotaExhaustedRef.current = false; setQuotaState({ userId, model: id, value: { status: "loading" } }); const conversationId = activeConversationIdRef.current; if (conversationId) { commitConversations((current) => current.map((conversation) => conversation.id === conversationId ? { ...conversation, modelId: id, updatedAt: new Date().toISOString() } : conversation, )); } const persistPreference = async () => { if ( !modelPreferenceMountedRef.current || preferenceEpoch !== modelPreferenceEpochRef.current ) return; const controller = new AbortController(); modelPreferenceControllerRef.current = controller; try { const response = await authFetch("/api/me/model", { method: "PATCH", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ model: id }), signal: controller.signal, }); if ( !modelPreferenceMountedRef.current || preferenceEpoch !== modelPreferenceEpochRef.current || controller.signal.aborted ) return; if (response.status === 401) onUnauthorizedRef.current?.(); if ( preferenceGeneration !== modelPreferenceGenerationRef.current || id !== activeModelIdRef.current ) return; if (response.ok) { void refreshQuota(id); return; } } catch { if ( !modelPreferenceMountedRef.current || preferenceEpoch !== modelPreferenceEpochRef.current || controller.signal.aborted || preferenceGeneration !== modelPreferenceGenerationRef.current || id !== activeModelIdRef.current ) return; } finally { if (modelPreferenceControllerRef.current === controller) { modelPreferenceControllerRef.current = null; } } quotaGenerationRef.current += 1; quotaControllerRef.current?.abort(); quotaExhaustedRef.current = false; setQuotaState({ userId, model: id, value: { status: "unavailable" } }); }; modelPreferenceChainRef.current = modelPreferenceChainRef.current.then( persistPreference, persistPreference, ); }, [commitConversations, refreshQuota, userId]); const sendMessage = useCallback(async (content: string, attachments: Attachment[] = []) => { const normalized = content.trim(); if ((!normalized && attachments.length === 0) || normalized.length > MAX_CHAT_MESSAGE_CHARACTERS || isGeneratingRef.current) return; await requestAssistant({ conversationId: activeConversationIdRef.current ?? undefined, model: activeModelIdRef.current, content: normalized || undefined, ...(attachments.length ? { attachmentIds: attachments.map((attachment) => attachment.id) } : {}), }, attachments); }, [requestAssistant]); const stopGenerating = useCallback(() => { abortControllerRef.current?.abort(); }, []); const retryLast = useCallback(async () => { const conversation = conversationsRef.current.find((item) => item.id === activeConversationIdRef.current); if (!conversation || !canRetryLastMessage(conversation)) return; await requestAssistant({ conversationId: conversation.id, model: conversation.modelId, retry: true }); }, [requestAssistant]); const activeConversation = conversations.find((conversation) => conversation.id === activeConversationId) ?? null; const quota = quotaState.userId === userId && quotaState.model === activeModelId ? quotaState.value : { status: "loading" as const }; return { conversations, activeConversation, activeModelId, isGenerating, canRetryLast: canRetryLastMessage(activeConversation), quota, selectConversation, newConversation, deleteConversation, setModel, sendMessage, stopGenerating, retryLast }; } function parseQuotaResponse(value: unknown, model: ModelId): QuotaSnapshot | null { if (!isRecord(value) || !isRecord(value.quota)) return null; const quota = value.quota; const availableTokens = quota.availableTokens; if ( quota.object !== "aegida.quota" || quota.model !== model || (availableTokens !== null && (typeof availableTokens !== "number" || !Number.isSafeInteger(availableTokens) || availableTokens < 0)) || typeof quota.unlimited !== "boolean" || typeof quota.exhausted !== "boolean" || (quota.resetsAt !== null && typeof quota.resetsAt !== "string") || typeof quota.asOf !== "string" ) { return null; } if (quota.unlimited) { if (availableTokens !== null || quota.exhausted || quota.resetsAt !== null) return null; } else if ( availableTokens === null || quota.exhausted !== (availableTokens === 0) ) { return null; } return quota as QuotaSnapshot; } function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } function canRetryLastMessage(conversation: Conversation | null): boolean { const last = conversation?.messages.at(-1); return Boolean( last?.role === "assistant" && ["error", "stopped"].includes(last.status) && !last.id.startsWith(CLIENT_MESSAGE_PREFIX), ); } function createTitle(content: string): string { const compact = content.replace(/\s+/g, " ").trim(); return compact.length <= 48 ? compact : `${compact.slice(0, 47).trimEnd()}…`; } function createOptimisticTurn(conversationId?: string): OptimisticTurn { return { conversationId: conversationId ?? `${CLIENT_CONVERSATION_PREFIX}${crypto.randomUUID()}`, userMessageId: `${CLIENT_MESSAGE_PREFIX}${crypto.randomUUID()}`, assistantMessageId: `${CLIENT_MESSAGE_PREFIX}${crypto.randomUUID()}`, createdAt: new Date().toISOString(), clientConversation: !conversationId, }; } function moveFirst(conversations: Conversation[], id: string) { const selected = conversations.find((conversation) => conversation.id === id); return selected ? [selected, ...conversations.filter((conversation) => conversation.id !== id)] : conversations; } async function readSafeError(response: Response): Promise { try { const value: unknown = await response.json(); if (typeof value === "object" && value !== null && "error" in value && typeof value.error === "string" && value.error.trim()) return value.error; } catch { // The API may not return JSON while it is unavailable. } return GENERIC_ERROR; }