aegida-console / hooks / use-chat.ts
use-chat.ts
Raw
"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<void>;
  stopGenerating(): void;
  retryLast(): Promise<void>;
};

export function useChat(
  userId: string,
  onUnauthorized?: () => void,
  initialModelId?: ModelId,
): UseChatResult {
  const [conversations, setConversations] = useState<Conversation[]>([]);
  const [activeConversationId, setActiveConversationId] = useState<string | null>(null);
  const [activeModelId, setActiveModelId] = useState<ModelId>(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<AbortController | null>(null);
  const streamingMessageRef = useRef<{ conversationId: string; messageId: string } | null>(null);
  const onUnauthorizedRef = useRef(onUnauthorized);
  const quotaControllerRef = useRef<AbortController | null>(null);
  const quotaGenerationRef = useRef(0);
  const quotaExhaustedRef = useRef(false);
  const modelPreferenceGenerationRef = useRef(0);
  const modelPreferenceChainRef = useRef<Promise<void>>(Promise.resolve());
  const modelPreferenceControllerRef = useRef<AbortController | null>(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<string, unknown> {
  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<string> {
  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;
}