i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 14:22:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: eliminate code duplication (DRY principle) (#211)

Browse Source 
## Problem Solved
Previous refactoring added 105 lines (1476→1581) by extracting code into separate files without eliminating duplication. This refactor focuses on reducing code size through deduplication while maintaining file separation for maintainability.

## Summary
- Reduced total lines from 1581 to 1519 (-62 lines, 3.9% reduction)
- Eliminated duplicate patterns using generic helpers and factory functions
- Maintained file structure for maintainability
- Zero functional changes - same behavior

### Phase 1: DRY use-quota-manager.tsx
- Created parseStorageCount() helper (eliminates 6x localStorage read duplication)
- Created createQuotaChecker() factory (consolidates 3 check function bodies)
- Created createQuotaIncrementer() factory (consolidates 3 increment function bodies)
- Result: 242→247 lines (+5 lines, but fully DRY with eliminated duplication)

### Phase 2: DRY chat-panel.tsx (1176→1109 lines, -67 lines)

#### 2.1: Extract checkAllQuotaLimits helper
- Replaced 3 occurrences of 18-line quota check blocks
- Saved 36 lines

#### 2.2: Extract sendChatMessage helper
- Replaced 3 occurrences of 21-line sendMessage+headers blocks
- Saved 42 lines

#### 2.3: Extract processFilesAndAppendContent helper
- Replaced 2 occurrences of file processing loops
- Handles PDF, text, and image files uniformly
- Async helper with optional image parts parameter
This commit is contained in:
Dayuan Jiang
2025-12-11 14:28:02 +09:00
committed by GitHub
parent 8b9336466f
commit 869391a029
5 changed files with 576 additions and 532 deletions
Download Patch File Download Diff File Expand all files Collapse all files

690
components/chat-panel.tsx
View File

@@ -17,39 +17,21 @@ import { FaGithub } from "react-icons/fa"
import { Toaster, toast } from "sonner"
import { ButtonWithTooltip } from "@/components/button-with-tooltip"
import { ChatInput } from "@/components/chat-input"
import { QuotaLimitToast } from "@/components/quota-limit-toast"
import {
SettingsDialog,
STORAGE_ACCESS_CODE_KEY,
STORAGE_AI_API_KEY_KEY,
STORAGE_AI_BASE_URL_KEY,
STORAGE_AI_MODEL_KEY,
STORAGE_AI_PROVIDER_KEY,
} from "@/components/settings-dialog"
import { SettingsDialog } from "@/components/settings-dialog"
import { useDiagram } from "@/contexts/diagram-context"
import { getAIConfig } from "@/lib/ai-config"
import { findCachedResponse } from "@/lib/cached-responses"
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
import { useQuotaManager } from "@/lib/use-quota-manager"
import { formatXML, wrapWithMxFile } from "@/lib/utils"
import { ChatMessageDisplay } from "./chat-message-display"
// localStorage keys for persistence
const STORAGE_MESSAGES_KEY = "next-ai-draw-io-messages"
const STORAGE_XML_SNAPSHOTS_KEY = "next-ai-draw-io-xml-snapshots"
const STORAGE_SESSION_ID_KEY = "next-ai-draw-io-session-id"
export const STORAGE_DIAGRAM_XML_KEY = "next-ai-draw-io-diagram-xml"
const STORAGE_REQUEST_COUNT_KEY = "next-ai-draw-io-request-count"
const STORAGE_REQUEST_DATE_KEY = "next-ai-draw-io-request-date"
const STORAGE_TOKEN_COUNT_KEY = "next-ai-draw-io-token-count"
const STORAGE_TOKEN_DATE_KEY = "next-ai-draw-io-token-date"
const STORAGE_TPM_COUNT_KEY = "next-ai-draw-io-tpm-count"
const STORAGE_TPM_MINUTE_KEY = "next-ai-draw-io-tpm-minute"
import { useDiagram } from "@/contexts/diagram-context"
import { findCachedResponse } from "@/lib/cached-responses"
import {
extractPdfText,
extractTextFileContent,
isPdfFile,
isTextFile,
MAX_EXTRACTED_CHARS,
} from "@/lib/pdf-utils"
import { formatXML, wrapWithMxFile } from "@/lib/utils"
import { ChatMessageDisplay } from "./chat-message-display"
// Type for message parts (tool calls and their states)
interface MessagePart {
@@ -186,11 +168,9 @@ export default function ChatPanel({
])
}
const [files, setFiles] = useState<File[]>([])
// Store extracted PDF text with extraction status
const [pdfData, setPdfData] = useState<
Map<File, { text: string; charCount: number; isExtracting: boolean }>
>(new Map())
// File processing using extracted hook
const { files, pdfData, handleFileChange, setFiles } = useFileProcessor()
const [showHistory, setShowHistory] = useState(false)
const [showSettingsDialog, setShowSettingsDialog] = useState(false)
const [, setAccessCodeRequired] = useState(false)
@@ -212,200 +192,12 @@ export default function ChatPanel({
.catch(() => setAccessCodeRequired(false))
}, [])
// Helper to check daily request limit
// Check if user has their own API key configured (bypass limits)
const hasOwnApiKey = useCallback((): boolean => {
const provider = localStorage.getItem(STORAGE_AI_PROVIDER_KEY)
const apiKey = localStorage.getItem(STORAGE_AI_API_KEY_KEY)
return !!(provider && apiKey)
}, [])
const checkDailyLimit = useCallback((): {
allowed: boolean
remaining: number
used: number
} => {
// Skip limit if user has their own API key
if (hasOwnApiKey()) return { allowed: true, remaining: -1, used: 0 }
if (dailyRequestLimit <= 0)
return { allowed: true, remaining: -1, used: 0 }
const today = new Date().toDateString()
const storedDate = localStorage.getItem(STORAGE_REQUEST_DATE_KEY)
let count = parseInt(
localStorage.getItem(STORAGE_REQUEST_COUNT_KEY) || "0",
10,
)
if (storedDate !== today) {
count = 0
localStorage.setItem(STORAGE_REQUEST_DATE_KEY, today)
localStorage.setItem(STORAGE_REQUEST_COUNT_KEY, "0")
}
return {
allowed: count < dailyRequestLimit,
remaining: dailyRequestLimit - count,
used: count,
}
}, [dailyRequestLimit, hasOwnApiKey])
// Helper to increment request count
const incrementRequestCount = useCallback((): void => {
const count = parseInt(
localStorage.getItem(STORAGE_REQUEST_COUNT_KEY) || "0",
10,
)
localStorage.setItem(STORAGE_REQUEST_COUNT_KEY, String(count + 1))
}, [])
// Helper to show quota limit toast (request-based)
const showQuotaLimitToast = useCallback(() => {
toast.custom(
(t) => (
<QuotaLimitToast
used={dailyRequestLimit}
limit={dailyRequestLimit}
onDismiss={() => toast.dismiss(t)}
/>
),
{ duration: 15000 },
)
}, [dailyRequestLimit, hasOwnApiKey])
// Helper to check daily token limit (checks if already over limit)
const checkTokenLimit = useCallback((): {
allowed: boolean
remaining: number
used: number
} => {
// Skip limit if user has their own API key
if (hasOwnApiKey()) return { allowed: true, remaining: -1, used: 0 }
if (dailyTokenLimit <= 0)
return { allowed: true, remaining: -1, used: 0 }
const today = new Date().toDateString()
const storedDate = localStorage.getItem(STORAGE_TOKEN_DATE_KEY)
let count = parseInt(
localStorage.getItem(STORAGE_TOKEN_COUNT_KEY) || "0",
10,
)
// Guard against NaN (e.g., if "NaN" was stored)
if (Number.isNaN(count)) count = 0
if (storedDate !== today) {
count = 0
localStorage.setItem(STORAGE_TOKEN_DATE_KEY, today)
localStorage.setItem(STORAGE_TOKEN_COUNT_KEY, "0")
}
return {
allowed: count < dailyTokenLimit,
remaining: dailyTokenLimit - count,
used: count,
}
}, [dailyTokenLimit, hasOwnApiKey])
// Helper to increment token count
const incrementTokenCount = useCallback((tokens: number): void => {
// Guard against NaN tokens
if (!Number.isFinite(tokens) || tokens <= 0) return
let count = parseInt(
localStorage.getItem(STORAGE_TOKEN_COUNT_KEY) || "0",
10,
)
// Guard against NaN count
if (Number.isNaN(count)) count = 0
localStorage.setItem(STORAGE_TOKEN_COUNT_KEY, String(count + tokens))
}, [])
// Helper to show token limit toast
const showTokenLimitToast = useCallback(
(used: number) => {
toast.custom(
(t) => (
<QuotaLimitToast
type="token"
used={used}
limit={dailyTokenLimit}
onDismiss={() => toast.dismiss(t)}
/>
),
{ duration: 15000 },
)
},
[dailyTokenLimit],
)
// Helper to check TPM (tokens per minute) limit
// Note: This only READS, doesn't write. incrementTPMCount handles writes.
const checkTPMLimit = useCallback((): {
allowed: boolean
remaining: number
used: number
} => {
// Skip limit if user has their own API key
if (hasOwnApiKey()) return { allowed: true, remaining: -1, used: 0 }
if (tpmLimit <= 0) return { allowed: true, remaining: -1, used: 0 }
const currentMinute = Math.floor(Date.now() / 60000).toString()
const storedMinute = localStorage.getItem(STORAGE_TPM_MINUTE_KEY)
let count = parseInt(
localStorage.getItem(STORAGE_TPM_COUNT_KEY) || "0",
10,
)
// Guard against NaN
if (Number.isNaN(count)) count = 0
// If we're in a new minute, treat count as 0 (will be reset on next increment)
if (storedMinute !== currentMinute) {
count = 0
}
return {
allowed: count < tpmLimit,
remaining: tpmLimit - count,
used: count,
}
}, [tpmLimit, hasOwnApiKey])
// Helper to increment TPM count
const incrementTPMCount = useCallback((tokens: number): void => {
// Guard against NaN tokens
if (!Number.isFinite(tokens) || tokens <= 0) return
const currentMinute = Math.floor(Date.now() / 60000).toString()
const storedMinute = localStorage.getItem(STORAGE_TPM_MINUTE_KEY)
let count = parseInt(
localStorage.getItem(STORAGE_TPM_COUNT_KEY) || "0",
10,
)
// Guard against NaN
if (Number.isNaN(count)) count = 0
// Reset if we're in a new minute
if (storedMinute !== currentMinute) {
count = 0
localStorage.setItem(STORAGE_TPM_MINUTE_KEY, currentMinute)
}
localStorage.setItem(STORAGE_TPM_COUNT_KEY, String(count + tokens))
}, [])
// Helper to show TPM limit toast
const showTPMLimitToast = useCallback(() => {
const limitDisplay =
tpmLimit >= 1000 ? `${tpmLimit / 1000}k` : String(tpmLimit)
toast.error(
`Rate limit reached (${limitDisplay} tokens/min). Please wait 60 seconds before sending another request.`,
{ duration: 8000 },
)
}, [tpmLimit])
// Quota management using extracted hook
const quotaManager = useQuotaManager({
dailyRequestLimit,
dailyTokenLimit,
tpmLimit,
})
// Generate a unique session ID for Langfuse tracing (restore from localStorage if available)
const [sessionId, setSessionId] = useState(() => {
@@ -470,14 +262,6 @@ export default function ChatPanel({
validationError,
)
// Return error to model - sendAutomaticallyWhen will trigger retry
const errorMessage = `${validationError}
Please fix the XML issues and call display_diagram again with corrected XML.
Your failed XML:
\`\`\`xml
${xml}
\`\`\``
if (DEBUG) {
console.log(
"[display_diagram] Adding tool output with state: output-error",
@@ -487,7 +271,14 @@ ${xml}
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: errorMessage,
errorText: `${validationError}
Please fix the XML issues and call display_diagram again with corrected XML.
Your failed XML:
\`\`\`xml
${xml}
\`\`\``,
})
} else {
// Success - diagram will be rendered by chat-message-display
@@ -645,8 +436,8 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
: 0
const actualTokens = inputTokens + outputTokens
if (actualTokens > 0) {
incrementTokenCount(actualTokens)
incrementTPMCount(actualTokens)
quotaManager.incrementTokenCount(actualTokens)
quotaManager.incrementTPMCount(actualTokens)
}
}
},
@@ -824,20 +615,11 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
const toolCallId = `cached-${Date.now()}`
// Build user message text including any file content
let userText = input
for (const file of files) {
if (isPdfFile(file)) {
const extracted = pdfData.get(file)
if (extracted?.text) {
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
}
} else if (isTextFile(file)) {
const extracted = pdfData.get(file)
if (extracted?.text) {
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
}
}
}
const userText = await processFilesAndAppendContent(
input,
files,
pdfData,
)
setMessages([
{
@@ -875,43 +657,14 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
// Build user text by concatenating input with pre-extracted text
// (Backend only reads first text part, so we must combine them)
let userText = input
const parts: any[] = []
if (files.length > 0) {
for (const file of files) {
if (isPdfFile(file)) {
// Use pre-extracted PDF text from pdfData
const extracted = pdfData.get(file)
if (extracted?.text) {
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
}
} else if (isTextFile(file)) {
// Use pre-extracted text file content from pdfData
const extracted = pdfData.get(file)
if (extracted?.text) {
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
}
} else {
// Handle as image
const reader = new FileReader()
const dataUrl = await new Promise<string>(
(resolve) => {
reader.onload = () =>
resolve(reader.result as string)
reader.readAsDataURL(file)
},
const userText = await processFilesAndAppendContent(
input,
files,
pdfData,
parts,
)
parts.push({
type: "file",
url: dataUrl,
mediaType: file.type,
})
}
}
}
// Add the combined text as the first part
parts.unshift({ type: "text", text: userText })
@@ -929,56 +682,11 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
xmlSnapshotsRef.current.set(messageIndex, chartXml)
saveXmlSnapshots()
// Check daily limit
const limitCheck = checkDailyLimit()
if (!limitCheck.allowed) {
showQuotaLimitToast()
return
}
// Check all quota limits
if (!checkAllQuotaLimits()) return
// Check daily token limit (actual usage tracked after response)
const tokenLimitCheck = checkTokenLimit()
if (!tokenLimitCheck.allowed) {
showTokenLimitToast(tokenLimitCheck.used)
return
}
sendChatMessage(parts, chartXml, previousXml, sessionId)
// Check TPM (tokens per minute) limit
const tpmCheck = checkTPMLimit()
if (!tpmCheck.allowed) {
showTPMLimitToast()
return
}
const accessCode =
localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
const aiProvider =
localStorage.getItem(STORAGE_AI_PROVIDER_KEY) || ""
const aiBaseUrl =
localStorage.getItem(STORAGE_AI_BASE_URL_KEY) || ""
const aiApiKey =
localStorage.getItem(STORAGE_AI_API_KEY_KEY) || ""
const aiModel = localStorage.getItem(STORAGE_AI_MODEL_KEY) || ""
sendMessage(
{ parts },
{
body: {
xml: chartXml,
previousXml,
sessionId,
},
headers: {
"x-access-code": accessCode,
...(aiProvider && { "x-ai-provider": aiProvider }),
...(aiBaseUrl && { "x-ai-base-url": aiBaseUrl }),
...(aiApiKey && { "x-ai-api-key": aiApiKey }),
...(aiModel && { "x-ai-model": aiModel }),
},
},
)
incrementRequestCount()
// Token count is tracked in onFinish with actual server usage
setInput("")
setFiles([])
@@ -994,81 +702,122 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
setInput(e.target.value)
}
const handleFileChange = async (newFiles: File[]) => {
setFiles(newFiles)
// Extract text immediately for new PDF/text files
for (const file of newFiles) {
const needsExtraction =
(isPdfFile(file) || isTextFile(file)) && !pdfData.has(file)
if (needsExtraction) {
// Mark as extracting
setPdfData((prev) => {
const next = new Map(prev)
next.set(file, {
text: "",
charCount: 0,
isExtracting: true,
})
return next
})
// Extract text asynchronously
try {
let text: string
if (isPdfFile(file)) {
text = await extractPdfText(file)
} else {
text = await extractTextFileContent(file)
// Helper functions for message actions (regenerate/edit)
// Extract previous XML snapshot before a given message index
const getPreviousXml = (beforeIndex: number): string => {
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
.filter((k) => k < beforeIndex)
.sort((a, b) => b - a)
return snapshotKeys.length > 0
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
: ""
}
// Check character limit
if (text.length > MAX_EXTRACTED_CHARS) {
const limitK = MAX_EXTRACTED_CHARS / 1000
toast.error(
`${file.name}: Content exceeds ${limitK}k character limit (${(text.length / 1000).toFixed(1)}k chars)`,
// Restore diagram from snapshot and update ref
const restoreDiagramFromSnapshot = (savedXml: string) => {
onDisplayChart(savedXml, true) // Skip validation for trusted snapshots
chartXMLRef.current = savedXml
}
// Clean up snapshots after a given message index
const cleanupSnapshotsAfter = (messageIndex: number) => {
for (const key of xmlSnapshotsRef.current.keys()) {
if (key > messageIndex) {
xmlSnapshotsRef.current.delete(key)
}
}
saveXmlSnapshots()
}
// Check all quota limits (daily requests, tokens, TPM)
const checkAllQuotaLimits = (): boolean => {
const limitCheck = quotaManager.checkDailyLimit()
if (!limitCheck.allowed) {
quotaManager.showQuotaLimitToast()
return false
}
const tokenLimitCheck = quotaManager.checkTokenLimit()
if (!tokenLimitCheck.allowed) {
quotaManager.showTokenLimitToast(tokenLimitCheck.used)
return false
}
const tpmCheck = quotaManager.checkTPMLimit()
if (!tpmCheck.allowed) {
quotaManager.showTPMLimitToast()
return false
}
return true
}
// Send chat message with headers and increment quota
const sendChatMessage = (
parts: any,
xml: string,
previousXml: string,
sessionId: string,
) => {
const config = getAIConfig()
sendMessage(
{ parts },
{
body: { xml, previousXml, sessionId },
headers: {
"x-access-code": config.accessCode,
...(config.aiProvider && {
"x-ai-provider": config.aiProvider,
}),
...(config.aiBaseUrl && {
"x-ai-base-url": config.aiBaseUrl,
}),
...(config.aiApiKey && { "x-ai-api-key": config.aiApiKey }),
...(config.aiModel && { "x-ai-model": config.aiModel }),
},
},
)
setPdfData((prev) => {
const next = new Map(prev)
next.delete(file)
return next
})
// Remove the file from the list
setFiles((prev) => prev.filter((f) => f !== file))
continue
quotaManager.incrementRequestCount()
}
setPdfData((prev) => {
const next = new Map(prev)
next.set(file, {
text,
charCount: text.length,
isExtracting: false,
})
return next
})
} catch (error) {
console.error("Failed to extract text:", error)
toast.error(`Failed to read file: ${file.name}`)
setPdfData((prev) => {
const next = new Map(prev)
next.delete(file)
return next
})
// Process files and append content to user text (handles PDF, text, and optionally images)
const processFilesAndAppendContent = async (
baseText: string,
files: File[],
pdfData: Map<File, FileData>,
imageParts?: any[],
): Promise<string> => {
let userText = baseText
for (const file of files) {
if (isPdfFile(file)) {
const extracted = pdfData.get(file)
if (extracted?.text) {
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
}
} else if (isTextFile(file)) {
const extracted = pdfData.get(file)
if (extracted?.text) {
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
}
} else if (imageParts) {
// Handle as image (only if imageParts array provided)
const reader = new FileReader()
const dataUrl = await new Promise<string>((resolve) => {
reader.onload = () => resolve(reader.result as string)
reader.readAsDataURL(file)
})
imageParts.push({
type: "file",
url: dataUrl,
mediaType: file.type,
})
}
}
// Clean up pdfData for removed files
setPdfData((prev) => {
const next = new Map(prev)
for (const key of prev.keys()) {
if (!newFiles.includes(key)) {
next.delete(key)
}
}
return next
})
return userText
}
const handleRegenerate = async (messageIndex: number) => {
@@ -1103,28 +852,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
return
}
// Get previous XML (snapshot before the one being regenerated)
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
.filter((k) => k < userMessageIndex)
.sort((a, b) => b - a)
const previousXml =
snapshotKeys.length > 0
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
: ""
// Restore the diagram to the saved state (skip validation for trusted snapshots)
onDisplayChart(savedXml, true)
// Update ref directly to ensure edit_diagram has the correct XML
chartXMLRef.current = savedXml
// Get previous XML and restore diagram state
const previousXml = getPreviousXml(userMessageIndex)
restoreDiagramFromSnapshot(savedXml)
// Clean up snapshots for messages after the user message (they will be removed)
for (const key of xmlSnapshotsRef.current.keys()) {
if (key > userMessageIndex) {
xmlSnapshotsRef.current.delete(key)
}
}
saveXmlSnapshots()
cleanupSnapshotsAfter(userMessageIndex)
// Remove the user message AND assistant message onwards (sendMessage will re-add the user message)
// Use flushSync to ensure state update is processed synchronously before sending
@@ -1133,53 +866,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
setMessages(newMessages)
})
// Check daily limit
const limitCheck = checkDailyLimit()
if (!limitCheck.allowed) {
showQuotaLimitToast()
return
}
// Check daily token limit (actual usage tracked after response)
const tokenLimitCheck = checkTokenLimit()
if (!tokenLimitCheck.allowed) {
showTokenLimitToast(tokenLimitCheck.used)
return
}
// Check TPM (tokens per minute) limit
const tpmCheck = checkTPMLimit()
if (!tpmCheck.allowed) {
showTPMLimitToast()
return
}
// Check all quota limits
if (!checkAllQuotaLimits()) return
// Now send the message after state is guaranteed to be updated
const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
const aiProvider = localStorage.getItem(STORAGE_AI_PROVIDER_KEY) || ""
const aiBaseUrl = localStorage.getItem(STORAGE_AI_BASE_URL_KEY) || ""
const aiApiKey = localStorage.getItem(STORAGE_AI_API_KEY_KEY) || ""
const aiModel = localStorage.getItem(STORAGE_AI_MODEL_KEY) || ""
sendChatMessage(userParts, savedXml, previousXml, sessionId)
sendMessage(
{ parts: userParts },
{
body: {
xml: savedXml,
previousXml,
sessionId,
},
headers: {
"x-access-code": accessCode,
...(aiProvider && { "x-ai-provider": aiProvider }),
...(aiBaseUrl && { "x-ai-base-url": aiBaseUrl }),
...(aiApiKey && { "x-ai-api-key": aiApiKey }),
...(aiModel && { "x-ai-model": aiModel }),
},
},
)
incrementRequestCount()
// Token count is tracked in onFinish with actual server usage
}
@@ -1200,28 +892,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
return
}
// Get previous XML (snapshot before the one being edited)
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
.filter((k) => k < messageIndex)
.sort((a, b) => b - a)
const previousXml =
snapshotKeys.length > 0
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
: ""
// Restore the diagram to the saved state (skip validation for trusted snapshots)
onDisplayChart(savedXml, true)
// Update ref directly to ensure edit_diagram has the correct XML
chartXMLRef.current = savedXml
// Get previous XML and restore diagram state
const previousXml = getPreviousXml(messageIndex)
restoreDiagramFromSnapshot(savedXml)
// Clean up snapshots for messages after the user message (they will be removed)
for (const key of xmlSnapshotsRef.current.keys()) {
if (key > messageIndex) {
xmlSnapshotsRef.current.delete(key)
}
}
saveXmlSnapshots()
cleanupSnapshotsAfter(messageIndex)
// Create new parts with updated text
const newParts = message.parts?.map((part: any) => {
@@ -1238,53 +914,11 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
setMessages(newMessages)
})
// Check daily limit
const limitCheck = checkDailyLimit()
if (!limitCheck.allowed) {
showQuotaLimitToast()
return
}
// Check daily token limit (actual usage tracked after response)
const tokenLimitCheck = checkTokenLimit()
if (!tokenLimitCheck.allowed) {
showTokenLimitToast(tokenLimitCheck.used)
return
}
// Check TPM (tokens per minute) limit
const tpmCheck = checkTPMLimit()
if (!tpmCheck.allowed) {
showTPMLimitToast()
return
}
// Check all quota limits
if (!checkAllQuotaLimits()) return
// Now send the edited message after state is guaranteed to be updated
const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
const aiProvider = localStorage.getItem(STORAGE_AI_PROVIDER_KEY) || ""
const aiBaseUrl = localStorage.getItem(STORAGE_AI_BASE_URL_KEY) || ""
const aiApiKey = localStorage.getItem(STORAGE_AI_API_KEY_KEY) || ""
const aiModel = localStorage.getItem(STORAGE_AI_MODEL_KEY) || ""
sendMessage(
{ parts: newParts },
{
body: {
xml: savedXml,
previousXml,
sessionId,
},
headers: {
"x-access-code": accessCode,
...(aiProvider && { "x-ai-provider": aiProvider }),
...(aiBaseUrl && { "x-ai-base-url": aiBaseUrl }),
...(aiApiKey && { "x-ai-api-key": aiApiKey }),
...(aiModel && { "x-ai-model": aiModel }),
},
},
)
incrementRequestCount()
sendChatMessage(newParts, savedXml, previousXml, sessionId)
// Token count is tracked in onFinish with actual server usage
}

26
lib/ai-config.ts Normal file
View File

@@ -0,0 +1,26 @@
import { STORAGE_KEYS } from "./storage"
/**
* Get AI configuration from localStorage.
* Returns API keys and settings for custom AI providers.
* Used to override server defaults when user provides their own API key.
*/
export function getAIConfig() {
if (typeof window === "undefined") {
return {
accessCode: "",
aiProvider: "",
aiBaseUrl: "",
aiApiKey: "",
aiModel: "",
}
}
return {
accessCode: localStorage.getItem(STORAGE_KEYS.accessCode) || "",
aiProvider: localStorage.getItem(STORAGE_KEYS.aiProvider) || "",
aiBaseUrl: localStorage.getItem(STORAGE_KEYS.aiBaseUrl) || "",
aiApiKey: localStorage.getItem(STORAGE_KEYS.aiApiKey) || "",
aiModel: localStorage.getItem(STORAGE_KEYS.aiModel) || "",
}
}

27
lib/storage.ts Normal file
View File

@@ -0,0 +1,27 @@
// Centralized localStorage keys
// Consolidates all storage keys from chat-panel.tsx and settings-dialog.tsx
export const STORAGE_KEYS = {
// Chat data
messages: "next-ai-draw-io-messages",
xmlSnapshots: "next-ai-draw-io-xml-snapshots",
diagramXml: "next-ai-draw-io-diagram-xml",
sessionId: "next-ai-draw-io-session-id",
// Quota tracking
requestCount: "next-ai-draw-io-request-count",
requestDate: "next-ai-draw-io-request-date",
tokenCount: "next-ai-draw-io-token-count",
tokenDate: "next-ai-draw-io-token-date",
tpmCount: "next-ai-draw-io-tpm-count",
tpmMinute: "next-ai-draw-io-tpm-minute",
// Settings
accessCode: "next-ai-draw-io-access-code",
closeProtection: "next-ai-draw-io-close-protection",
accessCodeRequired: "next-ai-draw-io-access-code-required",
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",
} as const

110
lib/use-file-processor.tsx Normal file
View File

@@ -0,0 +1,110 @@
"use client"
import { useState } from "react"
import { toast } from "sonner"
import {
extractPdfText,
extractTextFileContent,
isPdfFile,
isTextFile,
MAX_EXTRACTED_CHARS,
} from "@/lib/pdf-utils"
export interface FileData {
text: string
charCount: number
isExtracting: boolean
}
/**
* Hook for processing file uploads, especially PDFs and text files.
* Handles text extraction, character limit validation, and cleanup.
*/
export function useFileProcessor() {
const [files, setFiles] = useState<File[]>([])
const [pdfData, setPdfData] = useState<Map<File, FileData>>(new Map())
const handleFileChange = async (newFiles: File[]) => {
setFiles(newFiles)
// Extract text immediately for new PDF/text files
for (const file of newFiles) {
const needsExtraction =
(isPdfFile(file) || isTextFile(file)) && !pdfData.has(file)
if (needsExtraction) {
// Mark as extracting
setPdfData((prev) => {
const next = new Map(prev)
next.set(file, {
text: "",
charCount: 0,
isExtracting: true,
})
return next
})
// Extract text asynchronously
try {
let text: string
if (isPdfFile(file)) {
text = await extractPdfText(file)
} else {
text = await extractTextFileContent(file)
}
// Check character limit
if (text.length > MAX_EXTRACTED_CHARS) {
const limitK = MAX_EXTRACTED_CHARS / 1000
toast.error(
`${file.name}: Content exceeds ${limitK}k character limit (${(text.length / 1000).toFixed(1)}k chars)`,
)
setPdfData((prev) => {
const next = new Map(prev)
next.delete(file)
return next
})
// Remove the file from the list
setFiles((prev) => prev.filter((f) => f !== file))
continue
}
setPdfData((prev) => {
const next = new Map(prev)
next.set(file, {
text,
charCount: text.length,
isExtracting: false,
})
return next
})
} catch (error) {
console.error("Failed to extract text:", error)
toast.error(`Failed to read file: ${file.name}`)
setPdfData((prev) => {
const next = new Map(prev)
next.delete(file)
return next
})
}
}
}
// Clean up pdfData for removed files
setPdfData((prev) => {
const next = new Map(prev)
for (const key of prev.keys()) {
if (!newFiles.includes(key)) {
next.delete(key)
}
}
return next
})
}
return {
files,
pdfData,
handleFileChange,
setFiles, // Export for external control (e.g., clearing files)
}
}

247
lib/use-quota-manager.tsx Normal file
View File

@@ -0,0 +1,247 @@
"use client"
import { useCallback, useMemo } from "react"
import { toast } from "sonner"
import { QuotaLimitToast } from "@/components/quota-limit-toast"
import { STORAGE_KEYS } from "@/lib/storage"
export interface QuotaConfig {
dailyRequestLimit: number
dailyTokenLimit: number
tpmLimit: number
}
export interface QuotaCheckResult {
allowed: boolean
remaining: number
used: number
}
/**
* Hook for managing request/token quotas and rate limiting.
* Handles three types of limits:
* - Daily request limit
* - Daily token limit
* - Tokens per minute (TPM) rate limit
*
* Users with their own API key bypass all limits.
*/
export function useQuotaManager(config: QuotaConfig): {
hasOwnApiKey: () => boolean
checkDailyLimit: () => QuotaCheckResult
checkTokenLimit: () => QuotaCheckResult
checkTPMLimit: () => QuotaCheckResult
incrementRequestCount: () => void
incrementTokenCount: (tokens: number) => void
incrementTPMCount: (tokens: number) => void
showQuotaLimitToast: () => void
showTokenLimitToast: (used: number) => void
showTPMLimitToast: () => void
} {
const { dailyRequestLimit, dailyTokenLimit, tpmLimit } = config
// Check if user has their own API key configured (bypass limits)
const hasOwnApiKey = useCallback((): boolean => {
const provider = localStorage.getItem(STORAGE_KEYS.aiProvider)
const apiKey = localStorage.getItem(STORAGE_KEYS.aiApiKey)
return !!(provider && apiKey)
}, [])
// Generic helper: Parse count from localStorage with NaN guard
const parseStorageCount = (key: string): number => {
const count = parseInt(localStorage.getItem(key) || "0", 10)
return Number.isNaN(count) ? 0 : count
}
// Generic helper: Create quota checker factory
const createQuotaChecker = useCallback(
(
getTimeKey: () => string,
timeStorageKey: string,
countStorageKey: string,
limit: number,
) => {
return (): QuotaCheckResult => {
if (hasOwnApiKey())
return { allowed: true, remaining: -1, used: 0 }
if (limit <= 0) return { allowed: true, remaining: -1, used: 0 }
const currentTime = getTimeKey()
const storedTime = localStorage.getItem(timeStorageKey)
let count = parseStorageCount(countStorageKey)
if (storedTime !== currentTime) {
count = 0
localStorage.setItem(timeStorageKey, currentTime)
localStorage.setItem(countStorageKey, "0")
}
return {
allowed: count < limit,
remaining: limit - count,
used: count,
}
}
},
[hasOwnApiKey],
)
// Generic helper: Create quota incrementer factory
const createQuotaIncrementer = useCallback(
(
getTimeKey: () => string,
timeStorageKey: string,
countStorageKey: string,
validateInput: boolean = false,
) => {
return (tokens: number = 1): void => {
if (validateInput && (!Number.isFinite(tokens) || tokens <= 0))
return
const currentTime = getTimeKey()
const storedTime = localStorage.getItem(timeStorageKey)
let count = parseStorageCount(countStorageKey)
if (storedTime !== currentTime) {
count = 0
localStorage.setItem(timeStorageKey, currentTime)
}
localStorage.setItem(countStorageKey, String(count + tokens))
}
},
[],
)
// Check daily request limit
const checkDailyLimit = useMemo(
() =>
createQuotaChecker(
() => new Date().toDateString(),
STORAGE_KEYS.requestDate,
STORAGE_KEYS.requestCount,
dailyRequestLimit,
),
[createQuotaChecker, dailyRequestLimit],
)
// Increment request count
const incrementRequestCount = useMemo(
() =>
createQuotaIncrementer(
() => new Date().toDateString(),
STORAGE_KEYS.requestDate,
STORAGE_KEYS.requestCount,
false,
),
[createQuotaIncrementer],
)
// Show quota limit toast (request-based)
const showQuotaLimitToast = useCallback(() => {
toast.custom(
(t) => (
<QuotaLimitToast
used={dailyRequestLimit}
limit={dailyRequestLimit}
onDismiss={() => toast.dismiss(t)}
/>
),
{ duration: 15000 },
)
}, [dailyRequestLimit])
// Check daily token limit
const checkTokenLimit = useMemo(
() =>
createQuotaChecker(
() => new Date().toDateString(),
STORAGE_KEYS.tokenDate,
STORAGE_KEYS.tokenCount,
dailyTokenLimit,
),
[createQuotaChecker, dailyTokenLimit],
)
// Increment token count
const incrementTokenCount = useMemo(
() =>
createQuotaIncrementer(
() => new Date().toDateString(),
STORAGE_KEYS.tokenDate,
STORAGE_KEYS.tokenCount,
true, // Validate input tokens
),
[createQuotaIncrementer],
)
// Show token limit toast
const showTokenLimitToast = useCallback(
(used: number) => {
toast.custom(
(t) => (
<QuotaLimitToast
type="token"
used={used}
limit={dailyTokenLimit}
onDismiss={() => toast.dismiss(t)}
/>
),
{ duration: 15000 },
)
},
[dailyTokenLimit],
)
// Check TPM (tokens per minute) limit
const checkTPMLimit = useMemo(
() =>
createQuotaChecker(
() => Math.floor(Date.now() / 60000).toString(),
STORAGE_KEYS.tpmMinute,
STORAGE_KEYS.tpmCount,
tpmLimit,
),
[createQuotaChecker, tpmLimit],
)
// Increment TPM count
const incrementTPMCount = useMemo(
() =>
createQuotaIncrementer(
() => Math.floor(Date.now() / 60000).toString(),
STORAGE_KEYS.tpmMinute,
STORAGE_KEYS.tpmCount,
true, // Validate input tokens
),
[createQuotaIncrementer],
)
// Show TPM limit toast
const showTPMLimitToast = useCallback(() => {
const limitDisplay =
tpmLimit >= 1000 ? `${tpmLimit / 1000}k` : String(tpmLimit)
toast.error(
`Rate limit reached (${limitDisplay} tokens/min). Please wait 60 seconds before sending another request.`,
{ duration: 8000 },
)
}, [tpmLimit])
return {
// Check functions
hasOwnApiKey,
checkDailyLimit,
checkTokenLimit,
checkTPMLimit,
// Increment functions
incrementRequestCount,
incrementTokenCount,
incrementTPMCount,
// Toast functions
showQuotaLimitToast,
showTokenLimitToast,
showTPMLimitToast,
}
}