Hooks

Where it lives

• hooks/use-diagram-tool-handlers.tsGitHub ↗
  Read the script

  import type { MutableRefObject } from "react"
  import { useRef } from "react"
  import type { DiagramOperation } from "@/components/chat/types"
  import type {
      ValidationState,
      ValidationStatus,
  } from "@/components/chat/ValidationCard"
  import type { ValidationResult } from "@/lib/diagram-validator"
  import { formatValidationFeedback } from "@/lib/diagram-validator"
  import { isMxCellXmlComplete, wrapWithMxFile } from "@/lib/utils"
  
  const DEBUG = process.env.NODE_ENV === "development"
  
  interface ToolCall {
      toolCallId: string
      toolName: string
      input: unknown
  }
  
  type AddToolOutputSuccess = {
      tool: string
      toolCallId: string
      state?: "output-available"
      output: string
      errorText?: undefined
  }
  
  type AddToolOutputError = {
      tool: string
      toolCallId: string
      state: "output-error"
      output?: undefined
      errorText: string
  }
  
  type AddToolOutputParams = AddToolOutputSuccess | AddToolOutputError
  
  type AddToolOutputFn = (params: AddToolOutputParams) => void
  
  const MAX_VALIDATION_RETRIES = 3
  
  // Type for the validation function passed from useValidateDiagram hook
  type ValidateDiagramFn = (
      imageData: string,
      sessionId?: string,
  ) => Promise<ValidationResult>
  
  interface UseDiagramToolHandlersParams {
      partialXmlRef: MutableRefObject<string>
      editDiagramOriginalXmlRef: MutableRefObject<Map<string, string>>
      chartXMLRef: MutableRefObject<string>
      onDisplayChart: (xml: string, skipValidation?: boolean) => string | null
      onFetchChart: (saveToHistory?: boolean) => Promise<string>
      onExport: () => void
      captureValidationPng?: () => Promise<string | null>
      validateDiagram?: ValidateDiagramFn
      enableVlmValidation?: boolean
      sessionId?: string
      onValidationStateChange?: (
          toolCallId: string,
          state: ValidationState,
      ) => void
  }
  
  /**
   * Hook that creates the onToolCall handler for diagram-related tools.
   * Handles display_diagram, edit_diagram, and append_diagram tools.
   *
   * Note: addToolOutput is passed at call time (not hook init) because
   * it comes from useChat which creates a circular dependency.
   */
  export function useDiagramToolHandlers({
      partialXmlRef,
      editDiagramOriginalXmlRef,
      chartXMLRef,
      onDisplayChart,
      onFetchChart,
      onExport,
      captureValidationPng,
      validateDiagram,
      enableVlmValidation = true,
      sessionId,
      onValidationStateChange,
  }: UseDiagramToolHandlersParams) {
      // Track validation retry count per tool call
      const validationRetryCountRef = useRef<Map<string, number>>(new Map())
  
      // Helper to update validation state
      const updateValidationState = (
          toolCallId: string,
          status: ValidationStatus,
          options?: {
              attempt?: number
              maxAttempts?: number
              result?: ValidationResult
              error?: string
              imageData?: string
          },
      ) => {
          if (onValidationStateChange) {
              onValidationStateChange(toolCallId, {
                  status,
                  ...options,
              })
          }
      }
      const handleToolCall = async (
          { toolCall }: { toolCall: ToolCall },
          addToolOutput: AddToolOutputFn,
      ) => {
          if (DEBUG) {
              console.log(
                  `[onToolCall] Tool: ${toolCall.toolName}, CallId: ${toolCall.toolCallId}`,
              )
          }
  
          if (toolCall.toolName === "display_diagram") {
              await handleDisplayDiagram(toolCall, addToolOutput)
          } else if (toolCall.toolName === "edit_diagram") {
              await handleEditDiagram(toolCall, addToolOutput)
          } else if (toolCall.toolName === "append_diagram") {
              handleAppendDiagram(toolCall, addToolOutput)
          }
      }
  
      const handleDisplayDiagram = async (
          toolCall: ToolCall,
          addToolOutput: AddToolOutputFn,
      ) => {
          const { xml } = toolCall.input as { xml: string }
  
          // DEBUG: Log raw input to diagnose false truncation detection
          if (DEBUG) {
              console.log(
                  "[display_diagram] XML ending (last 100 chars):",
                  xml.slice(-100),
              )
              console.log("[display_diagram] XML length:", xml.length)
          }
  
          // Check if XML is truncated (incomplete mxCell indicates truncated output)
          const isTruncated = !isMxCellXmlComplete(xml)
          if (DEBUG) {
              console.log("[display_diagram] isTruncated:", isTruncated)
          }
  
          if (isTruncated) {
              // Store the partial XML for continuation via append_diagram
              partialXmlRef.current = xml
  
              // Tell LLM to use append_diagram to continue
              const partialEnding = partialXmlRef.current.slice(-500)
              addToolOutput({
                  tool: "display_diagram",
                  toolCallId: toolCall.toolCallId,
                  state: "output-error",
                  errorText: `Output was truncated due to length limits. Use the append_diagram tool to continue.
  
  Your output ended with:
  \`\`\`
  ${partialEnding}
  \`\`\`
  
  NEXT STEP: Call append_diagram with the continuation XML.
  - Do NOT include wrapper tags or root cells (id="0", id="1")
  - Start from EXACTLY where you stopped
  - Complete all remaining mxCell elements`,
              })
              return
          }
  
          // Complete XML received - use it directly
          // (continuation is now handled via append_diagram tool)
          const finalXml = xml
          partialXmlRef.current = "" // Reset any partial from previous truncation
  
          // Wrap raw XML with full mxfile structure for draw.io
          const fullXml = wrapWithMxFile(finalXml)
  
          // loadDiagram validates and returns error if invalid
          const validationError = onDisplayChart(fullXml)
  
          if (validationError) {
              console.warn("[display_diagram] Validation error:", validationError)
              // Return error to model - sendAutomaticallyWhen will trigger retry
              if (DEBUG) {
               

• hooks/use-dictionary.tsGitHub ↗
  Read the script

  "use client"
  
  import React, { createContext, useContext } from "react"
  import type { Dictionary } from "@/lib/i18n/dictionaries"
  
  const DictionaryContext = createContext<Dictionary | null>(null)
  
  export function DictionaryProvider({
      children,
      dictionary,
  }: React.PropsWithChildren<{ dictionary: Dictionary }>) {
      return React.createElement(
          DictionaryContext.Provider,
          { value: dictionary },
          children,
      )
  }
  
  export function useDictionary() {
      const dict = useContext(DictionaryContext)
      if (!dict) {
          throw new Error(
              "useDictionary must be used within a DictionaryProvider",
          )
      }
      return dict
  }
  
  export default useDictionary
  

• hooks/use-model-config.tsGitHub ↗
  Read the script

  "use client"
  
  import { useCallback, useEffect, useState } from "react"
  import { getApiEndpoint } from "@/lib/base-path"
  import type { FlattenedServerModel } from "@/lib/server-model-config"
  import { STORAGE_KEYS } from "@/lib/storage"
  import {
      createEmptyConfig,
      createModelConfig,
      createProviderConfig,
      type FlattenedModel,
      findModelById,
      flattenModels,
      type ModelConfig,
      type MultiModelConfig,
      type ProviderConfig,
      type ProviderName,
  } from "@/lib/types/model-config"
  
  // Old storage keys for migration
  const OLD_KEYS = {
      aiProvider: "next-ai-draw-io-ai-provider",
      aiBaseUrl: "next-ai-draw-io-ai-base-url",
      aiApiKey: "next-ai-draw-io-ai-api-key",
      aiModel: "next-ai-draw-io-ai-model",
  }
  
  /**
   * Migrate from old single-provider format to new multi-model format
   */
  function migrateOldConfig(): MultiModelConfig | null {
      if (typeof window === "undefined") return null
  
      const oldProvider = localStorage.getItem(OLD_KEYS.aiProvider)
      const oldApiKey = localStorage.getItem(OLD_KEYS.aiApiKey)
      const oldModel = localStorage.getItem(OLD_KEYS.aiModel)
  
      // No old config to migrate
      if (!oldProvider || !oldApiKey || !oldModel) return null
  
      const oldBaseUrl = localStorage.getItem(OLD_KEYS.aiBaseUrl)
  
      // Create new config from old format
      const provider = createProviderConfig(oldProvider as ProviderName)
      provider.apiKey = oldApiKey
      if (oldBaseUrl) provider.baseUrl = oldBaseUrl
  
      const model = createModelConfig(oldModel)
      provider.models.push(model)
  
      const config: MultiModelConfig = {
          version: 1,
          providers: [provider],
          selectedModelId: model.id,
      }
  
      // Clear old keys after migration
      localStorage.removeItem(OLD_KEYS.aiProvider)
      localStorage.removeItem(OLD_KEYS.aiBaseUrl)
      localStorage.removeItem(OLD_KEYS.aiApiKey)
      localStorage.removeItem(OLD_KEYS.aiModel)
  
      return config
  }
  
  /**
   * Load config from localStorage
   */
  function loadConfig(): MultiModelConfig {
      if (typeof window === "undefined") return createEmptyConfig()
  
      // First, check if new format exists
      const stored = localStorage.getItem(STORAGE_KEYS.modelConfigs)
      if (stored) {
          try {
              return JSON.parse(stored) as MultiModelConfig
          } catch {
              console.error("Failed to parse model config")
          }
      }
  
      // Try migration from old format
      const migrated = migrateOldConfig()
      if (migrated) {
          // Save migrated config
          localStorage.setItem(
              STORAGE_KEYS.modelConfigs,
              JSON.stringify(migrated),
          )
          return migrated
      }
  
      return createEmptyConfig()
  }
  
  /**
   * Save config to localStorage
   */
  function saveConfig(config: MultiModelConfig): void {
      if (typeof window === "undefined") return
      localStorage.setItem(STORAGE_KEYS.modelConfigs, JSON.stringify(config))
  }
  
  export interface UseModelConfigReturn {
      // State
      config: MultiModelConfig
      isLoaded: boolean
  
      // Getters
      models: FlattenedModel[]
      selectedModel: FlattenedModel | undefined
      selectedModelId: string | undefined
      showUnvalidatedModels: boolean
  
      // Actions
      setSelectedModelId: (modelId: string | undefined) => void
      setShowUnvalidatedModels: (show: boolean) => void
      addProvider: (provider: ProviderName) => ProviderConfig
      updateProvider: (
          providerId: string,
          updates: Partial<ProviderConfig>,
      ) => void
      deleteProvider: (providerId: string) => void
      addModel: (providerId: string, modelId: string) => ModelConfig
      updateModel: (
          providerId: string,
          modelConfigId: string,
          updates: Partial<ModelConfig>,
      ) => void
      deleteModel: (providerId: string, modelConfigId: string) => void
      resetConfig: () => void
  }
  
  export function useModelConfig(): UseModelConfigReturn {
      const [config, setConfig] = useState<MultiModelConfig>(createEmptyConfig)
      const [isLoaded, setIsLoaded] = useState(false)
      const [serverModels, setServerModels] = useState<FlattenedServerModel[]>([])
      const [serverLoaded, setServerLoaded] = useState(false)
  
      // Load client config on mount
      useEffect(() => {
          const loaded = loadConfig()
          setConfig(loaded)
          setIsLoaded(true)
      }, [])
  
      // Load server models on mount (if any)
      useEffect(() => {
          if (typeof window === "undefined") return
  
          fetch(getApiEndpoint("/api/server-models"))
              .then((res) => {
                  if (!res.ok) {
                      console.error(
                          "Failed to load server models:",
                          res.status,
                          res.statusText,
                      )
                      throw new Error(`Request failed with status ${res.status}`)
                  }
                  return res.json()
              })
              .then((data) => {
                  const raw: FlattenedServerModel[] = data?.models || []
                  setServerModels(raw)
                  setServerLoaded(true)
  
                  // Auto-select default server model if no model is currently selected
                  setConfig((prev) => {
                      if (!prev.selectedModelId && raw.length > 0) {
                          const defaultModel = raw.find((m) => m.isDefault)
                          if (defaultModel) {
                              return { ...prev, selectedModelId: defaultModel.id }
                          }
                          // If no default marked, use first server model
                          return { ...prev, selectedModelId: raw[0].id }
                      }
                      return prev
                  })
              })
              .catch((error) => {
                  console.error("Error while loading server models:", error)
                  setServerLoaded(true)
              })
      }, [])
  
      // Save config whenever it changes (after initial load)
      useEffect(() => {
          if (isLoaded) {
              saveConfig(confi

• hooks/use-session-manager.tsGitHub ↗
  Read the script

  "use client"
  
  import { useCallback, useEffect, useRef, useState } from "react"
  import {
      type ChatSession,
      createEmptySession,
      deleteSession as deleteSessionFromDB,
      enforceSessionLimit,
      extractTitle,
      getAllSessionMetadata,
      getSession,
      isIndexedDBAvailable,
      migrateFromLocalStorage,
      type SessionMetadata,
      type StoredMessage,
      saveSession,
  } from "@/lib/session-storage"
  
  export interface SessionData {
      messages: StoredMessage[]
      xmlSnapshots: [number, string][]
      diagramXml: string
      thumbnailDataUrl?: string
      diagramHistory?: { svg: string; xml: string }[]
  }
  
  export interface UseSessionManagerReturn {
      // State
      sessions: SessionMetadata[]
      currentSessionId: string | null
      currentSession: ChatSession | null
      isLoading: boolean
      isAvailable: boolean
  
      // Actions
      switchSession: (id: string) => Promise<SessionData | null>
      deleteSession: (id: string) => Promise<{ wasCurrentSession: boolean }>
      // forSessionId: optional session ID to verify save targets correct session (prevents stale debounce writes)
      saveCurrentSession: (
          data: SessionData,
          forSessionId?: string | null,
      ) => Promise<void>
      refreshSessions: () => Promise<void>
      clearCurrentSession: () => void
  }
  
  interface UseSessionManagerOptions {
      /** Session ID from URL param - if provided, load this session; if null, start blank */
      initialSessionId?: string | null
  }
  
  export function useSessionManager(
      options: UseSessionManagerOptions = {},
  ): UseSessionManagerReturn {
      const { initialSessionId } = options
      const [sessions, setSessions] = useState<SessionMetadata[]>([])
      const [currentSessionId, setCurrentSessionId] = useState<string | null>(
          null,
      )
      const [currentSession, setCurrentSession] = useState<ChatSession | null>(
          null,
      )
      const [isLoading, setIsLoading] = useState(true)
      const [isAvailable, setIsAvailable] = useState(false)
  
      const isInitializedRef = useRef(false)
      // Sequence guard for URL changes - prevents out-of-order async resolution
      const urlChangeSequenceRef = useRef(0)
  
      // Load sessions list
      const refreshSessions = useCallback(async () => {
          if (!isIndexedDBAvailable()) return
          try {
              const metadata = await getAllSessionMetadata()
              setSessions(metadata)
          } catch (error) {
              console.error("Failed to refresh sessions:", error)
          }
      }, [])
  
      // Initialize on mount
      useEffect(() => {
          if (isInitializedRef.current) return
          isInitializedRef.current = true
  
          async function init() {
              setIsLoading(true)
  
              if (!isIndexedDBAvailable()) {
                  setIsAvailable(false)
                  setIsLoading(false)
                  return
              }
  
              setIsAvailable(true)
  
              try {
                  // Run migration first (one-time conversion from localStorage)
                  await migrateFromLocalStorage()
  
                  // Load sessions list
                  const metadata = await getAllSessionMetadata()
                  setSessions(metadata)
  
                  // Only load a session if initialSessionId is provided (from URL param)
                  if (initialSessionId) {
                      const session = await getSession(initialSessionId)
                      if (session) {
                          setCurrentSession(session)
                          setCurrentSessionId(session.id)
                      }
                      // If session not found, stay in blank state (URL has invalid session ID)
                  }
                  // If no initialSessionId, start with blank state (no auto-restore)
              } catch (error) {
                  console.error("Failed to initialize session manager:", error)
              } finally {
                  setIsLoading(false)
              }
          }
  
          init()
      }, [initialSessionId])
  
      // Handle URL session ID changes after initialization
      // Note: intentionally NOT including currentSessionId in deps to avoid race conditions
      // when clearCurrentSession() is called before URL updates
      useEffect(() => {
          if (!isInitializedRef.current) return // Wait for initial load
          if (!isAvailable) return
  
          // Increment sequence to invalidate any pending async operations
          urlChangeSequenceRef.current++
          const currentSequence = urlChangeSequenceRef.current
  
          async function handleSessionIdChange() {
              if (initialSessionId) {
                  // URL has session ID - load it
                  const session = await getSession(initialSessionId)
  
                  // Check if this request is still the latest (sequence guard)
                  // If not, a newer URL change happened while we were loading
                  if (currentSequence !== urlChangeSequenceRef.current) {
                      return
                  }
  
                  if (session) {
                      // Only update if the session is different from current
                      setCurrentSessionId((current) => {
                          if (current !== session.id) {
                              setCurrentSession(session)
                              return session.id
                          }
                          return current
                      })
                  }
              }
              // Removed: else clause that clears session
              // Clearing is now handled explicitly by clearCurrentSession()
              // This prevents race conditions when URL update is async
          }
  
          handleSessionIdChange()
      }, [initialSessionId, isAvailable])
  
      // Refresh sessions on window focus (multi-tab sync)
      useEffect(() => {
          const handleFocus = () => {
              refreshSessions()
          }
          window.addEventListener("focus", handleFocus)
          return () => window.removeEventListener("focus", handleFocus)
      }, [refreshSessions])
  
      // Switch to a d

• hooks/use-validate-diagram.tsGitHub ↗
  Read the script

  "use client"
  
  /**
   * Hook for VLM-based diagram validation using AI SDK's useObject.
   */
  
  import { experimental_useObject as useObject } from "@ai-sdk/react"
  import { useCallback, useRef } from "react"
  import { getApiEndpoint } from "@/lib/base-path"
  import {
      type ValidationResult,
      ValidationResultSchema,
  } from "@/lib/validation-schema"
  
  export type { ValidationResult }
  
  // Default valid result for fallback cases
  const DEFAULT_VALID_RESULT: ValidationResult = {
      valid: true,
      issues: [],
      suggestions: [],
  }
  
  interface UseValidateDiagramOptions {
      onSuccess?: (result: ValidationResult) => void
      onError?: (error: Error) => void
  }
  
  // Track pending validation promises for imperative API
  type PendingValidation = {
      resolve: (result: ValidationResult) => void
      reject: (error: Error) => void
  }
  
  export function useValidateDiagram(options: UseValidateDiagramOptions = {}) {
      const { onSuccess, onError } = options
      const pendingValidationRef = useRef<PendingValidation | null>(null)
  
      const { object, submit, isLoading, error, stop } = useObject({
          api: getApiEndpoint("/api/validate-diagram"),
          schema: ValidationResultSchema,
          onFinish: ({
              object,
              error: finishError,
          }: {
              object: ValidationResult | undefined
              error: Error | undefined
          }) => {
              if (finishError) {
                  console.error(
                      "[useValidateDiagram] Validation error:",
                      finishError,
                  )
                  onError?.(finishError)
                  pendingValidationRef.current?.reject(finishError)
                  pendingValidationRef.current = null
                  return
              }
  
              if (object) {
                  const result = object as ValidationResult
                  onSuccess?.(result)
                  pendingValidationRef.current?.resolve(result)
                  pendingValidationRef.current = null
              }
          },
          onError: (err: Error) => {
              console.error("[useValidateDiagram] Stream error:", err)
              onError?.(err)
              pendingValidationRef.current?.reject(err)
              pendingValidationRef.current = null
          },
      })
  
      /**
       * Validate a diagram image.
       * Returns a promise that resolves with the validation result.
       */
      const validate = useCallback(
          async (
              imageData: string,
              sessionId?: string,
          ): Promise<ValidationResult> => {
              // Reject any pending validation to prevent promise leaks
              if (pendingValidationRef.current) {
                  pendingValidationRef.current.reject(
                      new Error("Validation superseded by new request"),
                  )
                  pendingValidationRef.current = null
              }
  
              return new Promise((resolve, reject) => {
                  // Store the promise handlers
                  pendingValidationRef.current = { resolve, reject }
  
                  // Submit the validation request
                  submit({ imageData, sessionId })
              })
          },
          [submit],
      )
  
      /**
       * Validate with fallback - returns default valid result on error.
       * Use this to avoid blocking the user on validation failures.
       */
      const validateWithFallback = useCallback(
          async (
              imageData: string,
              sessionId?: string,
          ): Promise<ValidationResult> => {
              try {
                  return await validate(imageData, sessionId)
              } catch (error) {
                  console.warn(
                      "[useValidateDiagram] Validation failed, using fallback:",
                      error,
                  )
                  return DEFAULT_VALID_RESULT
              }
          },
          [validate],
      )
  
      return {
          // Validation functions
          validate,
          validateWithFallback,
          stop,
  
          // State
          isValidating: isLoading,
          partialResult: object as ValidationResult | undefined,
          error,
      }
  }
  

Read the script before you install anything that runs on your machine. This is the one part of a plugin that acts without being asked.