snapshot of https://github.com/excalidraw/excalidraw.git (master, 2026-09-16; blobs > 131072B omitted)

packages/excalidraw/components/TTDDialog/hooks/useChatManagement.ts

@@ -0,0 +1,114 @@
+import { useCallback, useState } from "react";
+
+import { useAtom, useSetAtom } from "../../../editor-jotai";
+
+import { errorAtom, chatHistoryAtom } from "../TTDContext";
+
+import { useTTDChatStorage } from "../useTTDChatStorage";
+
+import { getLastAssistantMessage } from "../utils/chat";
+
+import type { SavedChat, TTDPersistenceAdapter } from "../types";
+
+interface UseChatManagementProps {
+  persistenceAdapter: TTDPersistenceAdapter;
+}
+
+export const useChatManagement = ({
+  persistenceAdapter,
+}: UseChatManagementProps) => {
+  const setError = useSetAtom(errorAtom);
+  const [chatHistory, setChatHistory] = useAtom(chatHistoryAtom);
+  const [isMenuOpen, setIsMenuOpen] = useState(false);
+
+  const { restoreChat, deleteChat, createNewChatId } = useTTDChatStorage({
+    persistenceAdapter,
+  });
+
+  const applyChatToState = useCallback(
+    (chat: SavedChat) => {
+      const restoredMessages = chat.messages.map((msg) => ({
+        ...msg,
+        timestamp:
+          msg.timestamp instanceof Date
+            ? msg.timestamp
+            : new Date(msg.timestamp),
+      }));
+
+      const history = {
+        id: chat.id,
+        messages: restoredMessages,
+        currentPrompt: "",
+      };
+
+      const lastAssistantMsg = getLastAssistantMessage(history);
+
+      setError(
+        lastAssistantMsg?.error ? new Error(lastAssistantMsg?.error) : null,
+      );
+      setChatHistory(history);
+    },
+    [setError, setChatHistory],
+  );
+
+  const resetChatState = useCallback(async () => {
+    const newSessionId = await createNewChatId();
+    setChatHistory({
+      id: newSessionId,
+      messages: [],
+      currentPrompt: "",
+    });
+    setError(null);
+  }, [createNewChatId, setChatHistory, setError]);
+
+  const onRestoreChat = useCallback(
+    (chat: SavedChat) => {
+      const restoredChat = restoreChat(chat);
+      applyChatToState(restoredChat);
+
+      setIsMenuOpen(false);
+    },
+    [restoreChat, applyChatToState],
+  );
+
+  const handleDeleteChat = useCallback(
+    async (chatId: string, event: React.MouseEvent) => {
+      event.stopPropagation();
+
+      const isDeletingActiveChat = chatId === chatHistory.id;
+      const updatedChats = await deleteChat(chatId);
+
+      if (isDeletingActiveChat) {
+        if (updatedChats.length > 0) {
+          const nextChat = updatedChats[0];
+          applyChatToState(nextChat);
+        } else {
+          await resetChatState();
+        }
+      }
+    },
+    [chatHistory.id, deleteChat, applyChatToState, resetChatState],
+  );
+
+  const handleNewChat = useCallback(async () => {
+    await resetChatState();
+    setIsMenuOpen(false);
+  }, [resetChatState]);
+
+  const handleMenuToggle = useCallback(() => {
+    setIsMenuOpen((prev) => !prev);
+  }, []);
+
+  const handleMenuClose = useCallback(() => {
+    setIsMenuOpen(false);
+  }, []);
+
+  return {
+    isMenuOpen,
+    onRestoreChat,
+    handleDeleteChat,
+    handleNewChat,
+    handleMenuToggle,
+    handleMenuClose,
+  };
+};

packages/excalidraw/components/TTDDialog/hooks/useMermaidRenderer.ts

@@ -0,0 +1,213 @@
+import { useCallback, useEffect, useMemo, useRef } from "react";
+
+import type { NonDeletedExcalidrawElement } from "@excalidraw/element/types";
+
+import { useAtom } from "../../../editor-jotai";
+
+import { chatHistoryAtom, errorAtom, showPreviewAtom } from "../TTDContext";
+import { convertMermaidToExcalidraw } from "../common";
+import { isValidMermaidSyntax } from "../utils/mermaidValidation";
+
+import { getLastAssistantMessage } from "../utils/chat";
+
+import { useUIAppState } from "../../../context/ui-appState";
+
+import type { BinaryFiles } from "../../../types";
+import type { MermaidToExcalidrawLibProps } from "../types";
+
+const FAST_THROTTLE_DELAY = 300;
+const SLOW_THROTTLE_DELAY = 3000;
+const RENDER_SPEED_THRESHOLD = 100;
+const PARSE_FAIL_DELAY = 100;
+
+interface UseMermaidRendererProps {
+  mermaidToExcalidrawLib: MermaidToExcalidrawLibProps;
+  canvasRef: React.RefObject<HTMLDivElement | null>;
+}
+
+export const useMermaidRenderer = ({
+  mermaidToExcalidrawLib,
+  canvasRef,
+}: UseMermaidRendererProps) => {
+  const [chatHistory] = useAtom(chatHistoryAtom);
+  const [, setError] = useAtom(errorAtom);
+
+  const [showPreview, setShowPreview] = useAtom(showPreviewAtom);
+  const isRenderingRef = useRef(false);
+
+  const lastAssistantMessage = useMemo(
+    () => getLastAssistantMessage(chatHistory),
+    [chatHistory],
+  );
+
+  // Keeping lastAssistantMesssage in ref, so I can access it in useEffect hooks
+  const lastAssistantMessageRef = useRef(lastAssistantMessage);
+  useEffect(() => {
+    lastAssistantMessageRef.current = lastAssistantMessage;
+  }, [lastAssistantMessage]);
+
+  const data = useRef<{
+    elements: readonly NonDeletedExcalidrawElement[];
+    files: BinaryFiles | null;
+  }>({
+    elements: [],
+    files: null,
+  });
+
+  const lastRenderTimeRef = useRef(0);
+  const pendingContentRef = useRef<string | null>(null);
+  const hasErrorOffsetRef = useRef(false);
+  const currentThrottleDelayRef = useRef(FAST_THROTTLE_DELAY);
+
+  const { theme } = useUIAppState();
+
+  const renderMermaid = useCallback(
+    async (mermaidDefinition: string): Promise<boolean> => {
+      if (!mermaidDefinition.trim() || !mermaidToExcalidrawLib.loaded) {
+        return false;
+      }
+
+      if (isRenderingRef.current) {
+        return false;
+      }
+
+      isRenderingRef.current = true;
+
+      const renderStartTime = performance.now();
+
+      const result = await convertMermaidToExcalidraw({
+        canvasRef,
+        data,
+        mermaidToExcalidrawLib,
+        setError,
+        mermaidDefinition,
+        theme,
+      });
+
+      const renderDuration = performance.now() - renderStartTime;
+
+      if (renderDuration < RENDER_SPEED_THRESHOLD) {
+        currentThrottleDelayRef.current = FAST_THROTTLE_DELAY;
+      } else {
+        currentThrottleDelayRef.current = SLOW_THROTTLE_DELAY;
+      }
+
+      isRenderingRef.current = false;
+      return result.success;
+    },
+    [canvasRef, mermaidToExcalidrawLib, setError, theme],
+  );
+
+  const throttledRenderMermaid = useMemo(() => {
+    const fn = async (content: string) => {
+      const now = Date.now();
+      const timeSinceLastRender = now - lastRenderTimeRef.current;
+      const throttleDelay = currentThrottleDelayRef.current;
+
+      if (!isValidMermaidSyntax(content)) {
+        if (!hasErrorOffsetRef.current) {
+          lastRenderTimeRef.current = Math.max(
+            lastRenderTimeRef.current,
+            now - throttleDelay + PARSE_FAIL_DELAY,
+          );
+          hasErrorOffsetRef.current = true;
+        }
+        pendingContentRef.current = content;
+        return;
+      }
+
+      hasErrorOffsetRef.current = false;
+
+      if (timeSinceLastRender < throttleDelay) {
+        pendingContentRef.current = content;
+        return;
+      }
+
+      pendingContentRef.current = null;
+      const success = await renderMermaid(content);
+      lastRenderTimeRef.current = Date.now();
+
+      if (!success) {
+        lastRenderTimeRef.current =
+          lastRenderTimeRef.current - throttleDelay + PARSE_FAIL_DELAY;
+        hasErrorOffsetRef.current = true;
+      }
+    };
+
+    fn.flush = async () => {
+      if (pendingContentRef.current) {
+        const content = pendingContentRef.current;
+        pendingContentRef.current = null;
+        await renderMermaid(content);
+        lastRenderTimeRef.current = Date.now();
+      }
+    };
+
+    fn.cancel = () => {
+      pendingContentRef.current = null;
+    };
+
+    return fn;
+  }, [renderMermaid]);
+
+  const resetThrottleState = useCallback(() => {
+    lastRenderTimeRef.current = 0;
+    pendingContentRef.current = null;
+    hasErrorOffsetRef.current = false;
+    currentThrottleDelayRef.current = FAST_THROTTLE_DELAY;
+  }, []);
+
+  // this hook is responsible for keep rendering during streaming
+  useEffect(() => {
+    if (lastAssistantMessage?.content && lastAssistantMessage?.isGenerating) {
+      throttledRenderMermaid(lastAssistantMessage.content);
+    } else if (!lastAssistantMessage?.isGenerating) {
+      throttledRenderMermaid.flush();
+      resetThrottleState();
+      if (lastAssistantMessage?.content) {
+        throttledRenderMermaid(lastAssistantMessage.content);
+      }
+    }
+  }, [
+    resetThrottleState,
+    throttledRenderMermaid,
+    lastAssistantMessage?.isGenerating,
+    lastAssistantMessage?.content,
+  ]);
+
+  // render the last message if the user navigates between the existing chats
+  useEffect(() => {
+    const msg = lastAssistantMessageRef.current;
+    if (!msg?.content || msg.error) {
+      return;
+    }
+
+    if (!showPreview) {
+      return;
+    }
+
+    renderMermaid(msg.content);
+  }, [chatHistory?.id, renderMermaid, showPreview]);
+
+  useEffect(() => {
+    if (
+      !chatHistory.messages?.filter((msg) => msg.type === "assistant").length
+    ) {
+      const canvasNode = canvasRef.current;
+      if (canvasNode) {
+        const parent = canvasNode.parentElement;
+        if (parent) {
+          parent.style.background = "";
+          canvasNode.replaceChildren();
+        }
+      }
+      setShowPreview(false);
+    } else if (!showPreview) {
+      setShowPreview(true);
+    }
+  }, [chatHistory.messages, setShowPreview, canvasRef, showPreview]);
+
+  return {
+    data,
+  };
+};

packages/excalidraw/components/TTDDialog/hooks/useTextGeneration.ts

@@ -0,0 +1,225 @@
+import { useRef } from "react";
+import { parseMermaidToExcalidraw } from "@excalidraw/mermaid-to-excalidraw";
+import { isFiniteNumber } from "@excalidraw/math";
+
+import { useAtom } from "../../../editor-jotai";
+
+import { trackEvent } from "../../../analytics";
+import { t } from "../../../i18n";
+
+import { errorAtom, rateLimitsAtom, chatHistoryAtom } from "../TTDContext";
+import { useChatAgent } from "../Chat";
+
+import {
+  addMessages,
+  getLastAssistantMessage,
+  getMessagesForLLM,
+  removeLastAssistantMessage,
+  updateAssistantContent,
+} from "../utils/chat";
+
+import type { LLMMessage, TTTDDialog } from "../types";
+
+const MIN_PROMPT_LENGTH = 3;
+const MAX_PROMPT_LENGTH = 10000;
+
+export const useTextGeneration = ({
+  onTextSubmit,
+}: {
+  onTextSubmit: (
+    props: TTTDDialog.OnTextSubmitProps,
+  ) => Promise<TTTDDialog.OnTextSubmitRetValue>;
+}) => {
+  const [, setError] = useAtom(errorAtom);
+  const [rateLimits, setRateLimits] = useAtom(rateLimitsAtom);
+  const [chatHistory, setChatHistory] = useAtom(chatHistoryAtom);
+
+  const { addUserMessage, addAssistantMessage, setAssistantError } =
+    useChatAgent();
+
+  const streamingAbortControllerRef = useRef<AbortController | null>(null);
+
+  const validatePrompt = (prompt: string): boolean => {
+    if (
+      prompt.length > MAX_PROMPT_LENGTH ||
+      prompt.length < MIN_PROMPT_LENGTH ||
+      rateLimits?.rateLimitRemaining === 0
+    ) {
+      if (prompt.length < MIN_PROMPT_LENGTH) {
+        setError(
+          new Error(
+            t("chat.errors.promptTooShort", { min: MIN_PROMPT_LENGTH }),
+          ),
+        );
+      }
+      if (prompt.length > MAX_PROMPT_LENGTH) {
+        setError(
+          new Error(t("chat.errors.promptTooLong", { max: MAX_PROMPT_LENGTH })),
+        );
+      }
+
+      return false;
+    }
+    return true;
+  };
+
+  const onGenerate: TTTDDialog.OnGenerate = async ({
+    prompt,
+    isRepairFlow = false,
+  }) => {
+    if (!validatePrompt(prompt)) {
+      return;
+    }
+
+    if (streamingAbortControllerRef.current) {
+      streamingAbortControllerRef.current.abort();
+    }
+
+    setError(null);
+
+    const abortController = new AbortController();
+    streamingAbortControllerRef.current = abortController;
+
+    if (!isRepairFlow) {
+      addUserMessage(prompt);
+      addAssistantMessage();
+    } else {
+      setChatHistory((prev) =>
+        updateAssistantContent(prev, {
+          isGenerating: true,
+          content: "",
+          error: undefined,
+          errorType: undefined,
+          errorDetails: undefined,
+        }),
+      );
+    }
+
+    try {
+      trackEvent("ai", "generate", "ttd");
+
+      const previousMessages = getMessagesForLLM(chatHistory);
+
+      const messages: LLMMessage[] = [
+        ...previousMessages.slice(-3),
+        { role: "user", content: prompt },
+      ];
+
+      const { generatedResponse, error, rateLimit, rateLimitRemaining } =
+        await onTextSubmit({
+          messages,
+          onStreamCreated: () => {
+            if (isRepairFlow) {
+              setChatHistory((prev) =>
+                updateAssistantContent(prev, {
+                  content: "",
+                  error: "",
+                  isGenerating: true,
+                }),
+              );
+            }
+          },
+          onChunk: (chunk: string) => {
+            setChatHistory((prev) => {
+              const lastAssistantMessage = getLastAssistantMessage(prev);
+              return updateAssistantContent(prev, {
+                content: lastAssistantMessage.content + chunk,
+              });
+            });
+          },
+          signal: abortController.signal,
+        });
+
+      setChatHistory((prev) =>
+        updateAssistantContent(prev, {
+          isGenerating: false,
+        }),
+      );
+
+      if (isFiniteNumber(rateLimit) && isFiniteNumber(rateLimitRemaining)) {
+        setRateLimits({ rateLimit, rateLimitRemaining });
+      }
+
+      if (error?.status === 429 || rateLimitRemaining === 0) {
+        setChatHistory((chatHistory) => {
+          if (error?.status === 429) {
+            chatHistory = removeLastAssistantMessage(chatHistory);
+          }
+
+          chatHistory = {
+            ...chatHistory,
+            messages: chatHistory.messages.filter(
+              (msg) =>
+                msg.type !== "warning" ||
+                msg.warningType === "rateLimitExceeded" ||
+                msg.warningType === "messageLimitExceeded",
+            ),
+          };
+          const messages = addMessages(chatHistory, [
+            {
+              type: "warning",
+              warningType:
+                rateLimitRemaining === 0
+                  ? "messageLimitExceeded"
+                  : "rateLimitExceeded",
+            },
+          ]);
+          return messages;
+        });
+      }
+
+      if (error) {
+        const isAborted =
+          error.name === "AbortError" ||
+          error.message === "Aborted" ||
+          abortController.signal.aborted;
+
+        // do nothing if the request was aborted by the user
+        if (isAborted) {
+          return;
+        }
+
+        const _error = new Error(
+          error.message || t("chat.errors.requestFailed"),
+        );
+        if (error.status !== 429) {
+          setAssistantError(_error.message, "network");
+        }
+        setError(_error);
+
+        return;
+      }
+
+      try {
+        await parseMermaidToExcalidraw(generatedResponse ?? "");
+        trackEvent("ai", "mermaid parse success", "ttd");
+      } catch (error: any) {
+        trackEvent("ai", "mermaid parse failed", "ttd");
+        const _error = new Error(
+          error.message || t("chat.errors.mermaidParseError"),
+        );
+        setAssistantError(_error.message, "parse");
+        setError(_error);
+      }
+    } catch (error: any) {
+      const _error = new Error(
+        error.message || t("chat.errors.generationFailed"),
+      );
+      setAssistantError(_error.message, "other");
+      setError(_error);
+    } finally {
+      streamingAbortControllerRef.current = null;
+    }
+  };
+
+  const handleAbort = () => {
+    if (streamingAbortControllerRef.current) {
+      streamingAbortControllerRef.current.abort();
+    }
+  };
+
+  return {
+    onGenerate,
+    handleAbort,
+  };
+};